Add df.columns property to improve visibility of generated API for …

…typed column access
Kotlin · Nov 19, 2024 · 6f522c6 · 6f522c6
1 parent 439f65d
commit 6f522c6
Show file tree

Hide file tree

Showing 15 changed files with 127 additions and 70 deletions.
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;
@@ -4198,6 +4202,7 @@ public final class org/jetbrains/kotlinx/dataframe/api/DataFrameGetKt {
 	public static final fun getColumnPath (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/columns/ColumnPath;
 	public static final fun getColumnPaths (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/util/List;
 	public static final fun getColumnWithPath (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/columns/ColumnWithPath;
+	public static final fun getColumns (Lorg/jetbrains/kotlinx/dataframe/DataFrame;)Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;
 	public static final fun getColumns (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/util/List;
 	public static final fun getColumns (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;)Ljava/util/List;
 	public static final fun getColumnsWithPaths (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/util/List;

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,36 @@ 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.columns. // easy to overview available columns
+ * ```
+ * In compiler plugin:
+ * ```
+ * val df = @Import DataFrame.read("file.csv")
+ * df.columns.
+ * ```
+ * By design, it can't be used in columns selection DSL. Because "columns" is optional - for convenient overview.
+ * No benefit in keeping it in the final code:
+ * ```
+ * // instead of this
+ * df.select { columns.myCol }
+ * // just write this
+ * df.select { myCol }
+ * ```
+ */
+public val <T> DataFrame<T>.columns: ColumnsScope<T> get() = 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/...ns/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/IrBodyFiller.kt b/...ns/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/.../kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TokenGenerator.kt b/.../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")