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

koperagen · koperagen · commit 1e1d07600e11 · 2024-12-12T18:21:27.000+02:00
diff --git a/core/api/core.api b/core/api/core.api
@@ -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
@@ -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;
@@ -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;
 }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/ColumnsContainer.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/ColumnsContainer.kt
@@ -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
 
@@ -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)
 
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/ColumnsScope.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/ColumnsScope.kt
@@ -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
+}
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/DataFrameGet.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/DataFrameGet.kt
@@ -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
@@ -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 =
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/codeGen/CodeGeneratorImpl.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/codeGen/CodeGeneratorImpl.kt
@@ -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
@@ -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!!
@@ -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")
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/CodeGenerationTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/CodeGenerationTests.kt
@@ -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
@@ -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!!
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/ReplCodeGenTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/ReplCodeGenTests.kt
@@ -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
@@ -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!!
@@ -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 {
@@ -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
     }
 }
diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/IrBodyFiller.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/IrBodyFiller.kt
@@ -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"))
     }
@@ -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
                 }
diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TokenGenerator.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TokenGenerator.kt
@@ -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
                         ),
diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/Names.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/Names.kt
@@ -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
diff --git a/plugins/kotlin-dataframe/testData/box/propertiesOrder.fir.txt b/plugins/kotlin-dataframe/testData/box/propertiesOrder.fir.txt
@@ -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|
diff --git a/plugins/kotlin-dataframe/testData/diagnostics/selectDuringTyping.fir.txt b/plugins/kotlin-dataframe/testData/diagnostics/selectDuringTyping.fir.txt
@@ -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|
diff --git a/plugins/kotlin-dataframe/testData/diagnostics/toDataFrame_java.fir.txt b/plugins/kotlin-dataframe/testData/diagnostics/toDataFrame_java.fir.txt
@@ -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|
@@ -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|
diff --git a/plugins/symbol-processor/src/main/kotlin/org/jetbrains/dataframe/ksp/ExtensionsGenerator.kt b/plugins/symbol-processor/src/main/kotlin/org/jetbrains/dataframe/ksp/ExtensionsGenerator.kt
@@ -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")
diff --git a/plugins/symbol-processor/src/test/kotlin/org/jetbrains/dataframe/ksp/DataFrameSymbolProcessorTest.kt b/plugins/symbol-processor/src/test/kotlin/org/jetbrains/dataframe/ksp/DataFrameSymbolProcessorTest.kt

Original file line number	Diff line number	Diff line change
`@@ -18,9 +18,9 @@ import kotlin.reflect.KProperty`
`18`	`18`	`*`
`19`	`19`	`* Base interface for [DataFrame] and [ColumnSelectionDsl]`
`20`	`20`	`*`
`21`		`- * @param T Schema marker. Used to generate extension properties for typed column access.`
	`21`	`+ * @param T Schema marker. Used to resolve generated extension properties for typed column access.`
`22`	`22`	`*/`
`23`		`-public interface ColumnsContainer<out T> {`
	`23`	`+public interface ColumnsContainer<out T> : ColumnsScope<T> {`
`24`	`24`
`25`	`25`	`// region columns`
`26`	`26`
`@@ -54,7 +54,7 @@ public interface ColumnsContainer<out T> {`
`54`	`54`
`55`	`55`	`// region get`
`56`	`56`
`57`		`- public operator fun get(columnName: String): AnyCol = getColumn(columnName)`
	`57`	`+ public override operator fun get(columnName: String): AnyCol = getColumn(columnName)`
`58`	`58`
`59`	`59`	`public operator fun get(columnPath: ColumnPath): AnyCol = getColumn(columnPath)`
`60`	`60`