init research

dataframe/plugins/kotlin-dataframe/README.md

@@ -0,0 +1,11 @@
+## ~~:plugins:kotlin-dataframe~~
+
+The Kotlin 2.x Compiler plugin of DataFrame.
+A plugin for your Kotlin project that can generate on-the-fly column accessors for the compiler and IDE even without
+having to provide data schemas!
+
+### DISABLED!
+
+Development of this module was moved to the Kotlin repository:
+https://github.com/JetBrains/kotlin/tree/master/plugins/kotlin-dataframe.
+These files are out of date.

dataframe/plugins/kotlin-dataframe/build.gradle.kts

@@ -0,0 +1,121 @@
+import org.jetbrains.kotlin.gradle.dsl.JvmTarget
+import org.jetbrains.kotlin.gradle.dsl.KotlinVersion
+
+plugins {
+    id("java")
+    kotlin("jvm")
+    kotlin("plugin.serialization")
+}
+
+group = "org.jetbrains.kotlinx.dataframe"
+
+val kotlinVersion: String by project.properties
+
+repositories {
+    mavenCentral()
+    maven("https://maven.pkg.jetbrains.space/kotlin/p/kotlin/dev/")
+}
+
+sourceSets {
+    main {
+        java.setSrcDirs(listOf("src"))
+        resources.setSrcDirs(listOf("resources"))
+    }
+    test {
+        java.setSrcDirs(listOf("tests", "tests-gen"))
+        resources.setSrcDirs(listOf("testResources"))
+    }
+}
+
+dependencies {
+    "org.jetbrains.kotlin:kotlin-compiler:$kotlinVersion".let {
+        compileOnly(it)
+        testImplementation(it)
+    }
+
+    testRuntimeOnly("org.jetbrains.kotlin:kotlin-test:$kotlinVersion")
+    testRuntimeOnly("org.jetbrains.kotlin:kotlin-script-runtime:$kotlinVersion")
+    testRuntimeOnly("org.jetbrains.kotlin:kotlin-annotations-jvm:$kotlinVersion")
+
+    implementation(project(projects.dataframeCompilerPluginCore.path, "shadow"))
+    testRuntimeOnly(projects.core)
+    testRuntimeOnly(projects.dataframeCsv)
+    testImplementation("org.jetbrains.kotlin:kotlin-reflect:$kotlinVersion")
+    testImplementation("org.jetbrains.kotlin:kotlin-compiler-internal-test-framework:$kotlinVersion")
+
+    testImplementation(platform("org.junit:junit-bom:5.11.3"))
+    testImplementation("org.junit.jupiter:junit-jupiter")
+    testImplementation("org.junit.platform:junit-platform-commons")
+    testImplementation("org.junit.platform:junit-platform-launcher")
+    testImplementation("org.junit.platform:junit-platform-runner")
+    testImplementation("org.junit.platform:junit-platform-suite-api")
+}
+
+tasks.test {
+    useJUnitPlatform()
+    jvmArgs("-Xmx2G")
+    environment("TEST_RESOURCES", project.layout.projectDirectory)
+    doFirst {
+        setLibraryProperty("org.jetbrains.kotlin.test.kotlin-stdlib", "kotlin-stdlib")
+        setLibraryProperty("org.jetbrains.kotlin.test.kotlin-reflect", "kotlin-reflect")
+        setLibraryProperty("org.jetbrains.kotlin.test.kotlin-test", "kotlin-test")
+        setLibraryProperty("org.jetbrains.kotlin.test.kotlin-script-runtime", "kotlin-script-runtime")
+        setLibraryProperty("org.jetbrains.kotlin.test.kotlin-annotations-jvm", "kotlin-annotations-jvm")
+    }
+}
+
+tasks.withType<org.jetbrains.kotlin.gradle.tasks.KotlinCompile>().configureEach {
+    friendPaths.from(project(projects.core.path).projectDir)
+    compilerOptions {
+        freeCompilerArgs.addAll(
+            "-Xcontext-receivers",
+        )
+        optIn.addAll(
+            "org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi",
+        )
+    }
+}
+
+tasks.withType<JavaCompile> {
+    sourceCompatibility = JavaVersion.VERSION_1_8.toString()
+    targetCompatibility = JavaVersion.VERSION_1_8.toString()
+}
+
+tasks.compileKotlin {
+    compilerOptions {
+        languageVersion = KotlinVersion.KOTLIN_2_0
+        jvmTarget = JvmTarget.JVM_1_8
+    }
+}
+
+tasks.compileTestKotlin {
+    compilerOptions {
+        languageVersion = KotlinVersion.KOTLIN_2_0
+        jvmTarget = JvmTarget.JVM_1_8
+    }
+}
+
+tasks.register<JavaExec>("generateTests") {
+    classpath = sourceSets.test.get().runtimeClasspath
+    mainClass = "org.jetbrains.kotlin.fir.dataframe.GenerateTestsKt"
+}
+
+fun Test.setLibraryProperty(propName: String, jarName: String) {
+    val path = project.configurations
+        .testRuntimeClasspath.get()
+        .files
+        .find { """$jarName-\d.*jar""".toRegex().matches(it.name) }
+        ?.absolutePath
+        ?: return
+    systemProperty(propName, path)
+}
+
+// Disabling all tests before removing the compiler plugin here
+// because we're moving to the Kotlin repo: #1290
+tasks.filter {
+    ":plugins:kotlin-dataframe" in it.path &&
+        "test" in it.name.lowercase()
+}.forEach {
+    println("disabling compiler plugin test task: ${it.path}. See #1290")
+    it.onlyIf { false }
+}

dataframe/plugins/kotlin-dataframe/gradle.properties

@@ -0,0 +1,2 @@
+kotlin.code.style=official
+kotlinVersion=2.0.20

dataframe/plugins/kotlin-dataframe/gradle/wrapper/gradle-wrapper.properties

@@ -0,0 +1,5 @@
+distributionBase=GRADLE_USER_HOME
+distributionPath=wrapper/dists
+distributionUrl=https\://services.gradle.org/distributions/gradle-8.10-bin.zip
+zipStoreBase=GRADLE_USER_HOME
+zipStorePath=wrapper/dists

dataframe/plugins/kotlin-dataframe/resources/META-INF/services/org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar

@@ -0,0 +1,6 @@
+#
+# Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+# Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+#
+
+org.jetbrains.kotlinx.dataframe.plugin.FirDataFrameComponentRegistrar

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/FirDataFrameComponentRegistrar.kt

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlinx.dataframe.plugin
+
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
+import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.extensions.FirExtensionApiInternals
+import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
+import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.DataRowSchemaSupertype
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.ExpressionAnalysisAdditionalChecker
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.FunctionCallTransformer
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.IrBodyFiller
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.ReturnTypeBasedReceiverInjector
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.TokenGenerator
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.TopLevelExtensionsGenerator
+
+class FirDataFrameExtensionRegistrar(
+    val isTest: Boolean,
+    val dumpSchemas: Boolean,
+) : FirExtensionRegistrar() {
+    @OptIn(FirExtensionApiInternals::class)
+    override fun ExtensionRegistrarContext.configurePlugin() {
+        +::TopLevelExtensionsGenerator
+        +::ReturnTypeBasedReceiverInjector
+        +{ it: FirSession ->
+            FunctionCallTransformer(it, isTest)
+        }
+        +::TokenGenerator
+        +::DataRowSchemaSupertype
+        +{ it: FirSession ->
+            ExpressionAnalysisAdditionalChecker(it, isTest, dumpSchemas)
+        }
+    }
+}
+
+@OptIn(ExperimentalCompilerApi::class)
+class FirDataFrameComponentRegistrar : CompilerPluginRegistrar() {
+    override fun ExtensionStorage.registerExtensions(configuration: CompilerConfiguration) {
+        FirExtensionRegistrarAdapter.registerExtension(
+            FirDataFrameExtensionRegistrar(isTest = false, dumpSchemas = true)
+        )
+        IrGenerationExtension.registerExtension(IrBodyFiller())
+    }
+
+    override val supportsK2: Boolean = true
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/analyzeRefinedCallShape.kt

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlinx.dataframe.plugin
+
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.resolvedType
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
+
+internal inline fun <reified T> KotlinTypeFacade.analyzeRefinedCallShape(
+    call: FirFunctionCall,
+    expectedReturnType: ClassId,
+    reporter: InterpretationErrorReporter
+): CallResult<T>? {
+    val callReturnType = call.resolvedType
+    if (callReturnType.classId != expectedReturnType) return null
+    // rootMarker is expected to be a token generated by the plugin.
+    // it's implied by "refined call"
+    // thus ConeClassLikeType
+    val rootMarkers = callReturnType.typeArguments.filterIsInstance<ConeClassLikeType>()
+    if (rootMarkers.size != callReturnType.typeArguments.size) return null
+
+    val newSchema: T? = call.interpreterName(session)?.let { name ->
+        when (name) {
+            else -> name.load<Interpreter<*>>().let { processor ->
+                val dataFrameSchema = interpret(call, processor, reporter = reporter)
+                    .let {
+                        val value = it?.value
+                        if (value !is T) {
+                            if (!reporter.errorReported) {
+                                reporter.reportInterpretationError(call, "${processor::class} must return ${T::class}, but was $value")
+                            }
+                            null
+                        } else {
+                            value
+                        }
+                    }
+                dataFrameSchema
+            }
+        }
+    }
+
+    return CallResult(rootMarkers, newSchema)
+}
+
+data class CallResult<T>(val markers: List<ConeClassLikeType>, val result: T?)
+
+class RefinedArguments(val refinedArguments: List<RefinedArgument>) : List<RefinedArgument> by refinedArguments
+
+data class RefinedArgument(val name: Name, val expression: FirExpression) {
+
+    override fun toString(): String {
+        return "RefinedArgument(name=$name, expression=${expression})"
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/CallShapeAttribute.kt

@@ -0,0 +1,21 @@
+package org.jetbrains.kotlinx.dataframe.plugin.extensions
+
+import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.SchemaProperty
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataKey
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataRegistry
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+
+sealed interface CallShapeData {
+    class Schema(val columns: List<SchemaProperty>) : CallShapeData
+
+    class Scope(val columns: List<SchemaProperty>, val source: KtSourceElement?) : CallShapeData
+
+    class RefinedType(val scopes: List<FirRegularClassSymbol>) : CallShapeData
+}
+
+
+object CallShapeAttribute : FirDeclarationDataKey()
+
+var FirClass.callShapeData: CallShapeData? by FirDeclarationDataRegistry.data(CallShapeAttribute)

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/DataFramePlugin.kt

@@ -0,0 +1,5 @@
+package org.jetbrains.kotlinx.dataframe.plugin.extensions
+
+import org.jetbrains.kotlin.GeneratedDeclarationKey
+
+data object DataFramePlugin : GeneratedDeclarationKey()

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/DataRowSchemaSupertype.kt

@@ -0,0 +1,52 @@
+package org.jetbrains.kotlinx.dataframe.plugin.extensions
+
+import org.jetbrains.kotlin.descriptors.ClassKind
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.checkers.toClassLikeSymbol
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.extensions.AnnotationFqn
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.FirSupertypeGenerationExtension
+import org.jetbrains.kotlin.fir.extensions.predicate.LookupPredicate
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.constructClassLikeType
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
+
+class DataRowSchemaSupertype(session: FirSession) : FirSupertypeGenerationExtension(session) {
+    companion object {
+        private val PREDICATE = LookupPredicate.create {
+            annotated(AnnotationFqn(DataSchema::class.java.name))
+        }
+        private val dataRowSchema = ClassId(FqName("org.jetbrains.kotlinx.dataframe.api"), Name.identifier("DataRowSchema"))
+    }
+
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(PREDICATE)
+    }
+
+    override fun needTransformSupertypes(declaration: FirClassLikeDeclaration): Boolean {
+        return declaration is FirRegularClass
+            && declaration.classKind == ClassKind.CLASS
+            && session.predicateBasedProvider.matches(PREDICATE, declaration)
+    }
+
+
+    override fun computeAdditionalSupertypes(
+        classLikeDeclaration: FirClassLikeDeclaration,
+        resolvedSupertypes: List<FirResolvedTypeRef>,
+        typeResolver: TypeResolveService
+    ): List<FirResolvedTypeRef> {
+        if (resolvedSupertypes.any { it.toClassLikeSymbol(session)?.classId == dataRowSchema }) return emptyList()
+        return listOf(
+            buildResolvedTypeRef {
+                type = dataRowSchema.constructClassLikeType(emptyArray())
+            }
+        )
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ExpressionAnalysisAdditionalChecker.kt

@@ -0,0 +1,259 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlinx.dataframe.plugin.extensions
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.diagnostics.AbstractSourceElementPositioningStrategy
+import org.jetbrains.kotlin.diagnostics.DiagnosticFactory1DelegateProvider
+import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.diagnostics.KtDiagnosticFactory1
+import org.jetbrains.kotlin.diagnostics.Severity
+import org.jetbrains.kotlin.diagnostics.SourceElementPositioningStrategies
+import org.jetbrains.kotlin.diagnostics.error1
+import org.jetbrains.kotlin.diagnostics.reportOn
+import org.jetbrains.kotlin.diagnostics.warning1
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.checkers.MppCheckerKind
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirPropertyChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirSimpleFunctionChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirFunctionCallChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirPropertyAccessExpressionChecker
+import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.declarations.hasAnnotation
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
+import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
+import org.jetbrains.kotlin.fir.references.toResolvedCallableSymbol
+import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirTypeProjectionWithVariance
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.isSubtypeOf
+import org.jetbrains.kotlin.fir.types.renderReadable
+import org.jetbrains.kotlin.fir.types.resolvedType
+import org.jetbrains.kotlin.fir.types.toSymbol
+import org.jetbrains.kotlin.fir.types.type
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.flatten
+import org.jetbrains.kotlinx.dataframe.plugin.pluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+import org.jetbrains.kotlinx.dataframe.plugin.utils.isDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.utils.isDataRow
+import org.jetbrains.kotlinx.dataframe.plugin.utils.isGroupBy
+
+class ExpressionAnalysisAdditionalChecker(
+    session: FirSession,
+    isTest: Boolean,
+    dumpSchemas: Boolean
+) : FirAdditionalCheckersExtension(session) {
+    override val expressionCheckers: ExpressionCheckers = object : ExpressionCheckers() {
+        override val functionCallCheckers: Set<FirFunctionCallChecker> = setOfNotNull(
+            Checker(isTest), FunctionCallSchemaReporter.takeIf { dumpSchemas }
+        )
+        override val propertyAccessExpressionCheckers: Set<FirPropertyAccessExpressionChecker> = setOfNotNull(
+            PropertyAccessSchemaReporter.takeIf { dumpSchemas }
+        )
+    }
+    override val declarationCheckers: DeclarationCheckers = object : DeclarationCheckers() {
+        override val propertyCheckers: Set<FirPropertyChecker> = setOfNotNull(PropertySchemaReporter.takeIf { dumpSchemas })
+        override val simpleFunctionCheckers: Set<FirSimpleFunctionChecker> = setOfNotNull(FunctionDeclarationSchemaReporter.takeIf { dumpSchemas })
+    }
+}
+
+private class Checker(
+    val isTest: Boolean,
+) : FirFunctionCallChecker(mppKind = MppCheckerKind.Common) {
+    companion object {
+        val ERROR by error1<KtElement, String>(SourceElementPositioningStrategies.DEFAULT)
+        val CAST_ERROR by error1<KtElement, String>(SourceElementPositioningStrategies.CALL_ELEMENT_WITH_DOT)
+        val CAST_TARGET_WARNING by warning1<KtElement, String>(SourceElementPositioningStrategies.CALL_ELEMENT_WITH_DOT)
+        val CAST_ID = CallableId(FqName.fromSegments(listOf("org", "jetbrains", "kotlinx", "dataframe", "api")), Name.identifier("cast"))
+        val CHECK = ClassId(FqName("org.jetbrains.kotlinx.dataframe.annotations"), Name.identifier("Check"))
+    }
+
+    override fun check(expression: FirFunctionCall, context: CheckerContext, reporter: DiagnosticReporter) {
+        with(KotlinTypeFacadeImpl(context.session, isTest)) {
+            analyzeCast(expression, reporter, context)
+//            analyzeRefinedCallShape(expression, reporter = object : InterpretationErrorReporter {
+//                override var errorReported: Boolean = false
+//
+//                override fun reportInterpretationError(call: FirFunctionCall, message: String) {
+//                    reporter.reportOn(call.source, ERROR, message, context)
+//                    errorReported = true
+//                }
+//
+//                override fun doNotReportInterpretationError() {
+//                    errorReported = true
+//                }
+//            })
+        }
+    }
+
+    private fun KotlinTypeFacadeImpl.analyzeCast(expression: FirFunctionCall, reporter: DiagnosticReporter, context: CheckerContext) {
+        val calleeReference = expression.calleeReference
+        if (calleeReference !is FirResolvedNamedReference
+            || calleeReference.toResolvedCallableSymbol()?.callableId != CAST_ID
+            || !calleeReference.resolvedSymbol.hasAnnotation(CHECK, session)) {
+            return
+        }
+        val targetProjection = expression.typeArguments.getOrNull(0) as? FirTypeProjectionWithVariance ?: return
+        val targetType = targetProjection.typeRef.coneType as? ConeClassLikeType ?: return
+        val targetSymbol = targetType.toSymbol(session)
+        if (targetSymbol != null && !targetSymbol.hasAnnotation(Names.DATA_SCHEMA_CLASS_ID, session)) {
+            val text = "Annotate ${targetType.renderReadable()} with @DataSchema to use generated properties"
+            reporter.reportOn(expression.source, CAST_TARGET_WARNING, text, context)
+        }
+        val coneType = expression.explicitReceiver?.resolvedType
+        if (coneType != null) {
+            val sourceType = coneType.fullyExpandedType(session).typeArguments.getOrNull(0)?.type as? ConeClassLikeType
+                ?: return
+            val source = pluginDataFrameSchema(sourceType)
+            val target = pluginDataFrameSchema(targetType)
+            val sourceColumns = source.flatten(includeFrames = true)
+            val targetColumns = target.flatten(includeFrames = true)
+            val sourceMap = sourceColumns.associate { it.path.path to it.column }
+            val missingColumns = mutableListOf<String>()
+            var valid = true
+            for (target in targetColumns) {
+                val source = sourceMap[target.path.path]
+                val present = if (source != null) {
+                    if (source !is SimpleDataColumn || target.column !is SimpleDataColumn) { continue }
+                    if (source.type.type().isSubtypeOf(target.column.type.type(), session)) {
+                        true
+                    } else {
+                        missingColumns += "${target.path.path} ${target.column.name}: ${source.type.type().renderReadable()} is not subtype of ${target.column.type.type()}"
+                        false
+                    }
+                } else {
+                    missingColumns += "${target.path.path} ${target.column.name} is missing"
+                    false
+                }
+
+                valid = valid && present
+            }
+            if (!valid) {
+                reporter.reportOn(expression.source, CAST_ERROR, "Cast cannot succeed \n ${missingColumns.joinToString("\n")}", context)
+            }
+        }
+    }
+}
+
+private data object PropertySchemaReporter : FirPropertyChecker(mppKind = MppCheckerKind.Common) {
+    val SCHEMA by info1<KtElement, String>(SourceElementPositioningStrategies.DECLARATION_NAME)
+
+    override fun check(declaration: FirProperty, context: CheckerContext, reporter: DiagnosticReporter) {
+        context.sessionContext {
+            declaration.returnTypeRef.coneType.let { type ->
+                reportSchema(reporter, declaration.source, SCHEMA, type, context)
+            }
+        }
+    }
+}
+
+private data object FunctionCallSchemaReporter : FirFunctionCallChecker(mppKind = MppCheckerKind.Common) {
+    val SCHEMA by info1<KtElement, String>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
+
+    override fun check(expression: FirFunctionCall, context: CheckerContext, reporter: DiagnosticReporter) {
+        if (expression.calleeReference.name in setOf(Name.identifier("let"), Name.identifier("run"))) return
+        val initializer = expression.resolvedType
+        context.sessionContext {
+            reportSchema(reporter, expression.source, SCHEMA, initializer, context)
+        }
+    }
+}
+
+private data object PropertyAccessSchemaReporter : FirPropertyAccessExpressionChecker(mppKind = MppCheckerKind.Common) {
+    val SCHEMA by info1<KtElement, String>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
+
+    override fun check(
+        expression: FirPropertyAccessExpression,
+        context: CheckerContext,
+        reporter: DiagnosticReporter
+    ) {
+        val initializer = expression.resolvedType
+        context.sessionContext {
+            reportSchema(reporter, expression.source, SCHEMA, initializer, context)
+        }
+    }
+}
+
+private data object FunctionDeclarationSchemaReporter : FirSimpleFunctionChecker(mppKind = MppCheckerKind.Common) {
+    val SCHEMA by info1<KtElement, String>(SourceElementPositioningStrategies.DECLARATION_SIGNATURE)
+
+    override fun check(declaration: FirSimpleFunction, context: CheckerContext, reporter: DiagnosticReporter) {
+        val type = declaration.returnTypeRef.coneType
+        context.sessionContext {
+            reportSchema(reporter, declaration.source, SCHEMA, type, context)
+        }
+    }
+}
+
+private fun SessionContext.reportSchema(
+    reporter: DiagnosticReporter,
+    source: KtSourceElement?,
+    factory: KtDiagnosticFactory1<String>,
+    type: ConeKotlinType,
+    context: CheckerContext,
+) {
+    val expandedType = type.fullyExpandedType(session)
+    var schema: PluginDataFrameSchema? = null
+    when {
+        expandedType.isDataFrame(session) -> {
+            schema = expandedType.typeArguments.getOrNull(0)?.let {
+                pluginDataFrameSchema(it)
+            }
+        }
+
+        expandedType.isDataRow(session) -> {
+            schema = expandedType.typeArguments.getOrNull(0)?.let {
+                pluginDataFrameSchema(it)
+            }
+        }
+
+        expandedType.isGroupBy(session) -> {
+            val keys = expandedType.typeArguments.getOrNull(0)
+            val grouped = expandedType.typeArguments.getOrNull(1)
+            if (keys != null && grouped != null) {
+                val keysSchema = pluginDataFrameSchema(keys)
+                val groupedSchema = pluginDataFrameSchema(grouped)
+                schema = PluginDataFrameSchema(
+                    listOf(
+                        SimpleColumnGroup("keys", keysSchema.columns()),
+                        SimpleFrameColumn("groups", groupedSchema.columns())
+                    )
+                )
+            }
+        }
+    }
+    if (schema != null && source != null) {
+        reporter.reportOn(source, factory, "\n" + schema.toString(), context)
+    }
+}
+
+fun CheckerContext.sessionContext(f: SessionContext.() -> Unit) {
+    SessionContext(session).f()
+}
+
+inline fun <reified P : PsiElement, A> info1(
+    positioningStrategy: AbstractSourceElementPositioningStrategy = SourceElementPositioningStrategies.DEFAULT
+): DiagnosticFactory1DelegateProvider<A> {
+    return DiagnosticFactory1DelegateProvider(Severity.INFO, positioningStrategy, P::class)
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt

@@ -0,0 +1,601 @@
+package org.jetbrains.kotlinx.dataframe.plugin.extensions
+
+import org.jetbrains.kotlin.cli.common.repl.replEscapeLineBreaks
+import org.jetbrains.kotlin.contracts.description.EventOccurrencesRange
+import org.jetbrains.kotlin.descriptors.ClassKind
+import org.jetbrains.kotlin.descriptors.EffectiveVisibility
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.FirAnnotationContainer
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirFunctionTarget
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.checkers.fullyExpandedClassId
+import org.jetbrains.kotlinx.dataframe.plugin.InterpretationErrorReporter
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.SchemaProperty
+import org.jetbrains.kotlinx.dataframe.plugin.analyzeRefinedCallShape
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+import org.jetbrains.kotlinx.dataframe.plugin.utils.projectOverDataColumnType
+import org.jetbrains.kotlin.fir.declarations.EmptyDeprecationsProvider
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.InlineStatus
+import org.jetbrains.kotlin.fir.declarations.builder.buildAnonymousFunction
+import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
+import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
+import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
+import org.jetbrains.kotlin.fir.declarations.utils.classId
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
+import org.jetbrains.kotlin.fir.expressions.buildResolvedArgumentList
+import org.jetbrains.kotlin.fir.expressions.builder.buildAnonymousFunctionExpression
+import org.jetbrains.kotlin.fir.expressions.builder.buildBlock
+import org.jetbrains.kotlin.fir.expressions.builder.buildFunctionCall
+import org.jetbrains.kotlin.fir.expressions.builder.buildPropertyAccessExpression
+import org.jetbrains.kotlin.fir.expressions.builder.buildReturnExpression
+import org.jetbrains.kotlin.fir.extensions.FirExtensionApiInternals
+import org.jetbrains.kotlin.fir.extensions.FirFunctionCallRefinementExtension
+import org.jetbrains.kotlin.fir.moduleData
+import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
+import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
+import org.jetbrains.kotlin.fir.resolve.calls.candidate.CallInfo
+import org.jetbrains.kotlin.fir.resolve.defaultType
+import org.jetbrains.kotlin.fir.resolve.fqName
+import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
+import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
+import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLookupTagWithFixedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirValueParameterSymbol
+import org.jetbrains.kotlin.fir.toFirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.ConeKotlinTypeProjection
+import org.jetbrains.kotlin.fir.types.ConeStarProjection
+import org.jetbrains.kotlin.fir.types.ConeTypeProjection
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.builder.buildTypeProjectionWithVariance
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.fir.types.impl.FirImplicitAnyTypeRef
+import org.jetbrains.kotlin.fir.types.resolvedType
+import org.jetbrains.kotlin.fir.types.toClassSymbol
+import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.text
+import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.PropertyName
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBy
+import kotlin.math.abs
+
+@OptIn(FirExtensionApiInternals::class)
+class FunctionCallTransformer(
+    session: FirSession,
+    override val isTest: Boolean,
+) : FirFunctionCallRefinementExtension(session), KotlinTypeFacade {
+    companion object {
+        const val DEFAULT_NAME = "DataFrameType"
+    }
+
+    private interface CallTransformer {
+        fun interceptOrNull(callInfo: CallInfo, symbol: FirNamedFunctionSymbol, hash: String): CallReturnType?
+
+        /**
+         * must still generate let with declared class from interceptOrNull when interpretation fails.
+         * it should only return null if later some frontend checker fails compilation in general
+         */
+        fun transformOrNull(call: FirFunctionCall, originalSymbol: FirNamedFunctionSymbol): FirFunctionCall?
+    }
+
+    // also update [ReturnTypeBasedReceiverInjector.SCHEMA_TYPES]
+    private val transformers = listOf(
+        GroupByCallTransformer(),
+        DataFrameCallTransformer(),
+        DataRowCallTransformer(),
+        ColumnGroupCallTransformer(),
+    )
+
+    override fun intercept(callInfo: CallInfo, symbol: FirNamedFunctionSymbol): CallReturnType? {
+        val callSiteAnnotations = (callInfo.callSite as? FirAnnotationContainer)?.annotations ?: emptyList()
+        if (callSiteAnnotations.any { it.fqName(session)?.shortName()?.equals(Name.identifier("DisableInterpretation")) == true }) {
+            return null
+        }
+        val noRefineAnnotation =
+            symbol.annotations.none { it.fqName(session)?.shortName()?.equals(Name.identifier("Refine")) == true }
+        val optIn = symbol.annotations.any { it.fqName(session)?.shortName()?.equals(Name.identifier("OptInRefine")) == true } &&
+            callSiteAnnotations.any { it.fqName(session)?.shortName()?.equals(Name.identifier("Import")) == true }
+        if (noRefineAnnotation && !optIn) {
+            return null
+        }
+        if (exposesLocalType(callInfo)) return null
+
+        val hash = run {
+            val hash = callInfo.name.hashCode() + callInfo.arguments.sumOf {
+                when (it) {
+                    is FirLiteralExpression -> it.value.hashCode()
+                    else -> it.source?.text?.hashCode() ?: 42
+                }
+            }
+            hashToTwoCharString(abs(hash))
+        }
+
+        return transformers.firstNotNullOfOrNull { it.interceptOrNull(callInfo, symbol, hash) }
+    }
+
+    private fun exposesLocalType(callInfo: CallInfo): Boolean {
+        val property = callInfo.containingDeclarations.lastOrNull()?.symbol as? FirPropertySymbol
+        return (property != null && !property.resolvedStatus.effectiveVisibility.privateApi)
+    }
+
+    private fun hashToTwoCharString(hash: Int): String {
+        val baseChars = "0123456789"
+        val base = baseChars.length
+        val positiveHash = abs(hash)
+        val char1 = baseChars[positiveHash % base]
+        val char2 = baseChars[(positiveHash / base) % base]
+
+        return "$char1$char2"
+    }
+
+    override fun transform(call: FirFunctionCall, originalSymbol: FirNamedFunctionSymbol): FirFunctionCall {
+        return transformers
+            .firstNotNullOfOrNull { it.transformOrNull(call, originalSymbol) }
+            ?: call
+    }
+
+    inner class DataSchemaLikeCallTransformer(val classId: ClassId) : CallTransformer {
+        override fun interceptOrNull(callInfo: CallInfo, symbol: FirNamedFunctionSymbol, hash: String): CallReturnType? {
+            if (symbol.resolvedReturnType.fullyExpandedClassId(session) != classId) return null
+            // possibly null if explicit receiver type is typealias
+            val argument = (callInfo.explicitReceiver?.resolvedType)?.typeArguments?.getOrNull(0)
+            val newDataFrameArgument = buildNewTypeArgument(argument, callInfo.name, hash)
+
+            val lookupTag = ConeClassLikeLookupTagImpl(classId)
+            val typeRef = buildResolvedTypeRef {
+                type = ConeClassLikeTypeImpl(
+                    lookupTag,
+                    arrayOf(
+                        ConeClassLikeTypeImpl(
+                            ConeClassLookupTagWithFixedSymbol(newDataFrameArgument.classId, newDataFrameArgument.symbol),
+                            emptyArray(),
+                            isNullable = false
+                        )
+                    ),
+                    isNullable = false
+                )
+            }
+            return CallReturnType(typeRef)
+        }
+
+        @OptIn(SymbolInternals::class)
+        override fun transformOrNull(call: FirFunctionCall, originalSymbol: FirNamedFunctionSymbol): FirFunctionCall? {
+            val callResult = analyzeRefinedCallShape<PluginDataFrameSchema>(call, classId, InterpretationErrorReporter.DEFAULT)
+            val (tokens, dataFrameSchema) = callResult ?: return null
+            val token = tokens[0]
+            val firstSchema = token.toClassSymbol(session)?.resolvedSuperTypes?.get(0)!!.toRegularClassSymbol(session)?.fir!!
+            val dataSchemaApis = materialize(dataFrameSchema ?: PluginDataFrameSchema.EMPTY, call, firstSchema)
+
+            val tokenFir = token.toClassSymbol(session)!!.fir
+            tokenFir.callShapeData = CallShapeData.RefinedType(dataSchemaApis.map { it.scope.symbol })
+
+            return buildScopeFunctionCall(call, originalSymbol, dataSchemaApis, listOf(tokenFir))
+        }
+    }
+
+    inner class DataFrameCallTransformer : CallTransformer by DataSchemaLikeCallTransformer(Names.DF_CLASS_ID)
+
+    inner class DataRowCallTransformer : CallTransformer by DataSchemaLikeCallTransformer(Names.DATA_ROW_CLASS_ID)
+
+    inner class ColumnGroupCallTransformer : CallTransformer by DataSchemaLikeCallTransformer(Names.COLUM_GROUP_CLASS_ID)
+
+    inner class GroupByCallTransformer : CallTransformer {
+        override fun interceptOrNull(
+            callInfo: CallInfo,
+            symbol: FirNamedFunctionSymbol,
+            hash: String
+        ): CallReturnType? {
+            if (symbol.resolvedReturnType.fullyExpandedClassId(session) != Names.GROUP_BY_CLASS_ID) return null
+            val keys = buildNewTypeArgument(null, Name.identifier("Key"), hash)
+            val group = buildNewTypeArgument(null, Name.identifier("Group"), hash)
+            val lookupTag = ConeClassLikeLookupTagImpl(Names.GROUP_BY_CLASS_ID)
+            val typeRef = buildResolvedTypeRef {
+                type = ConeClassLikeTypeImpl(
+                    lookupTag,
+                    arrayOf(
+                        ConeClassLikeTypeImpl(
+                            ConeClassLookupTagWithFixedSymbol(keys.classId, keys.symbol),
+                            emptyArray<ConeTypeProjection>(),
+                            isNullable = false
+                        ),
+                        ConeClassLikeTypeImpl(
+                            ConeClassLookupTagWithFixedSymbol(group.classId, group.symbol),
+                            emptyArray<ConeTypeProjection>(),
+                            isNullable = false
+                        )
+                    ),
+                    isNullable = false
+                )
+            }
+            return CallReturnType(typeRef)
+        }
+
+        @OptIn(SymbolInternals::class)
+        override fun transformOrNull(call: FirFunctionCall, originalSymbol: FirNamedFunctionSymbol): FirFunctionCall? {
+            val callResult = analyzeRefinedCallShape<GroupBy>(call, Names.GROUP_BY_CLASS_ID, InterpretationErrorReporter.DEFAULT)
+            val (rootMarkers, groupBy) = callResult ?: return null
+
+            val keyMarker = rootMarkers[0]
+            val groupMarker = rootMarkers[1]
+
+            val (keySchema, groupSchema) = if (groupBy != null) {
+                val keySchema = groupBy.keys
+                val groupSchema = groupBy.groups
+                keySchema to groupSchema
+            } else {
+                PluginDataFrameSchema.EMPTY to PluginDataFrameSchema.EMPTY
+            }
+
+            val firstSchema = keyMarker.toClassSymbol(session)?.resolvedSuperTypes?.get(0)!!.toRegularClassSymbol(session)?.fir!!
+            val firstSchema1 = groupMarker.toClassSymbol(session)?.resolvedSuperTypes?.get(0)!!.toRegularClassSymbol(session)?.fir!!
+
+            val keyApis = materialize(keySchema, call, firstSchema, "Key")
+            val groupApis = materialize(groupSchema, call, firstSchema1, "Group", i = keyApis.size)
+
+            val groupToken = keyMarker.toClassSymbol(session)!!.fir
+            groupToken.callShapeData = CallShapeData.RefinedType(keyApis.map { it.scope.symbol })
+
+            val keyToken = groupMarker.toClassSymbol(session)!!.fir
+            keyToken.callShapeData = CallShapeData.RefinedType(groupApis.map { it.scope.symbol })
+
+            return buildScopeFunctionCall(call, originalSymbol, keyApis + groupApis, additionalDeclarations = listOf(groupToken, keyToken))
+        }
+    }
+
+    private fun buildNewTypeArgument(argument: ConeTypeProjection?, name: Name, hash: String): FirRegularClass {
+        val suggestedName = if (argument == null) {
+            "${name.asTokenName()}_$hash"
+        } else {
+            when (argument) {
+                is ConeStarProjection -> {
+                    "${name.asTokenName()}_$hash"
+                }
+                is ConeKotlinTypeProjection -> {
+                    val titleCase = argument.type.classId?.shortClassName
+                        ?.identifierOrNullIfSpecial?.titleCase()
+                        ?.substringBeforeLast("_")
+                        ?: DEFAULT_NAME
+                    "${titleCase}_$hash"
+                }
+            }
+        }
+        val tokenId = nextName("${suggestedName}I")
+        val token = buildSchema(tokenId)
+
+        val dataFrameTypeId = nextName(suggestedName)
+        val dataFrameType = buildRegularClass {
+            moduleData = session.moduleData
+            resolvePhase = FirResolvePhase.BODY_RESOLVE
+            origin = FirDeclarationOrigin.Source
+            status = FirResolvedDeclarationStatusImpl(Visibilities.Local, Modality.ABSTRACT, EffectiveVisibility.Local)
+            deprecationsProvider = EmptyDeprecationsProvider
+            classKind = ClassKind.CLASS
+            scopeProvider = FirKotlinScopeProvider()
+            superTypeRefs += buildResolvedTypeRef {
+                type = ConeClassLikeTypeImpl(
+                    ConeClassLookupTagWithFixedSymbol(tokenId, token.symbol),
+                    emptyArray(),
+                    isNullable = false
+                )
+            }
+
+            this.name = dataFrameTypeId.shortClassName
+            this.symbol = FirRegularClassSymbol(dataFrameTypeId)
+        }
+        return dataFrameType
+    }
+
+    private fun nextName(s: String) = ClassId(CallableId.PACKAGE_FQ_NAME_FOR_LOCAL, FqName(s), true)
+
+    private fun Name.asTokenName() = identifierOrNullIfSpecial?.titleCase() ?: DEFAULT_NAME
+
+    @OptIn(SymbolInternals::class)
+    private fun buildScopeFunctionCall(
+        call: FirFunctionCall,
+        originalSymbol: FirNamedFunctionSymbol,
+        dataSchemaApis: List<DataSchemaApi>,
+        additionalDeclarations: List<FirClass>
+    ): FirFunctionCall {
+
+        val explicitReceiver = call.explicitReceiver
+        val receiverType = explicitReceiver?.resolvedType
+        val returnType = call.resolvedType
+        val scopeFunction = if (explicitReceiver != null) findLet() else findRun()
+        val originalSource = call.calleeReference.source
+
+        // original call is inserted later
+        call.transformCalleeReference(object : FirTransformer<Nothing?>() {
+            override fun <E : FirElement> transformElement(element: E, data: Nothing?): E {
+                return if (element is FirResolvedNamedReference) {
+                    @Suppress("UNCHECKED_CAST")
+                    buildResolvedNamedReference {
+                        this.name = element.name
+                        resolvedSymbol = originalSymbol
+                    } as E
+                } else {
+                    element
+                }
+            }
+        }, null)
+
+        val callExplicitReceiver = call.explicitReceiver
+        val callDispatchReceiver = call.dispatchReceiver
+        val callExtensionReceiver = call.extensionReceiver
+
+        val argument = buildAnonymousFunctionExpression {
+            isTrailingLambda = true
+            val fSymbol = FirAnonymousFunctionSymbol()
+            val target = FirFunctionTarget(null, isLambda = true)
+            anonymousFunction = buildAnonymousFunction {
+                resolvePhase = FirResolvePhase.BODY_RESOLVE
+                moduleData = session.moduleData
+                origin = FirDeclarationOrigin.Source
+                status = FirResolvedDeclarationStatusImpl(Visibilities.Local, Modality.FINAL, EffectiveVisibility.Local)
+                deprecationsProvider = EmptyDeprecationsProvider
+                returnTypeRef = buildResolvedTypeRef {
+                    type = returnType
+                }
+                val parameterSymbol = receiverType?.let {
+                    val itName = Name.identifier("it")
+                    val parameterSymbol = FirValueParameterSymbol(itName)
+                    valueParameters += buildValueParameter {
+                        moduleData = session.moduleData
+                        origin = FirDeclarationOrigin.Source
+                        returnTypeRef = buildResolvedTypeRef {
+                            type = receiverType
+                        }
+                        this.name = itName
+                        this.symbol = parameterSymbol
+                        containingFunctionSymbol = fSymbol
+                        isCrossinline = false
+                        isNoinline = false
+                        isVararg = false
+                    }
+                    parameterSymbol
+                }
+                body = buildBlock {
+                    this.coneTypeOrNull = returnType
+                    dataSchemaApis.asReversed().forEach {
+                        statements += it.schema
+                        statements += it.scope
+                    }
+
+                    statements += additionalDeclarations
+
+                    statements += buildReturnExpression {
+                        if (parameterSymbol != null) {
+                            val itPropertyAccess = buildPropertyAccessExpression {
+                                coneTypeOrNull = receiverType
+                                calleeReference = buildResolvedNamedReference {
+                                    name = parameterSymbol.name
+                                    resolvedSymbol = parameterSymbol
+                                }
+                            }
+                            if (callDispatchReceiver != null) {
+                                call.replaceDispatchReceiver(itPropertyAccess)
+                            }
+                            call.replaceExplicitReceiver(itPropertyAccess)
+                            if (callExtensionReceiver != null) {
+                                call.replaceExtensionReceiver(itPropertyAccess)
+                            }
+                        }
+
+                        result = call
+                        this.target = target
+                    }
+                }
+                this.symbol = fSymbol
+                isLambda = true
+                hasExplicitParameterList = false
+                typeRef = buildResolvedTypeRef {
+                    type = if (receiverType != null) {
+                        ConeClassLikeTypeImpl(
+                            ConeClassLikeLookupTagImpl(ClassId(FqName("kotlin"), Name.identifier("Function1"))),
+                            typeArguments = arrayOf(receiverType, returnType),
+                            isNullable = false
+                        )
+                    } else {
+                        ConeClassLikeTypeImpl(
+                            ConeClassLikeLookupTagImpl(ClassId(FqName("kotlin"), Name.identifier("Function0"))),
+                            typeArguments = arrayOf(returnType),
+                            isNullable = false
+                        )
+                    }
+                }
+                invocationKind = EventOccurrencesRange.EXACTLY_ONCE
+                inlineStatus = InlineStatus.Inline
+            }.also {
+                target.bind(it)
+            }
+        }
+
+        val newCall1 = buildFunctionCall {
+            // source = call.source makes IDE navigate to `let` declaration
+            source = null
+            this.coneTypeOrNull = returnType
+            if (receiverType != null) {
+                typeArguments += buildTypeProjectionWithVariance {
+                    typeRef = buildResolvedTypeRef {
+                        type = receiverType
+                    }
+                    variance = Variance.INVARIANT
+                }
+            }
+
+            typeArguments += buildTypeProjectionWithVariance {
+                typeRef = buildResolvedTypeRef {
+                    type = returnType
+                }
+                variance = Variance.INVARIANT
+            }
+            dispatchReceiver = null
+            this.explicitReceiver = callExplicitReceiver
+            extensionReceiver = callExtensionReceiver ?: callDispatchReceiver
+            argumentList = buildResolvedArgumentList(
+                original = null,
+                linkedMapOf(argument to scopeFunction.valueParameterSymbols[0].fir)
+            )
+            calleeReference = buildResolvedNamedReference {
+                source = originalSource
+                this.name = scopeFunction.name
+                resolvedSymbol = scopeFunction
+            }
+        }
+        return newCall1
+    }
+
+    private fun materialize(
+        dataFrameSchema: PluginDataFrameSchema,
+        call: FirFunctionCall,
+        firstSchema: FirRegularClass,
+        prefix: String = "",
+        i: Int = 0
+    ): List<DataSchemaApi> {
+        var i = i
+        val dataSchemaApis = mutableListOf<DataSchemaApi>()
+        val usedNames = mutableMapOf<String, Int>()
+        fun PluginDataFrameSchema.materialize(
+            schema: FirRegularClass? = null,
+            suggestedName: String? = null
+        ): DataSchemaApi {
+            val schema = if (schema != null) {
+                schema
+            } else {
+                requireNotNull(suggestedName)
+                val uniqueSuffix = usedNames.compute(suggestedName) { _, i -> (i ?: 0) + 1 }
+                val name = nextName(suggestedName + uniqueSuffix)
+                buildSchema(name)
+            }
+
+            val scopeId = ClassId(CallableId.PACKAGE_FQ_NAME_FOR_LOCAL, FqName("Scope${i++}"), true)
+            val scope = buildRegularClass {
+                moduleData = session.moduleData
+                resolvePhase = FirResolvePhase.BODY_RESOLVE
+                origin = FirDeclarationOrigin.Source
+                status = FirResolvedDeclarationStatusImpl(Visibilities.Local, Modality.FINAL, EffectiveVisibility.Local)
+                deprecationsProvider = EmptyDeprecationsProvider
+                classKind = ClassKind.CLASS
+                scopeProvider = FirKotlinScopeProvider()
+                superTypeRefs += FirImplicitAnyTypeRef(null)
+
+                this.name = scopeId.shortClassName
+                this.symbol = FirRegularClassSymbol(scopeId)
+            }
+
+            val properties = columns().map {
+                fun PluginDataFrameSchema.materialize(column: SimpleCol): DataSchemaApi {
+                    val text = call.source?.text ?: call.calleeReference.name
+                    val name =
+                        "${column.name.titleCase().replEscapeLineBreaks()}_${hashToTwoCharString(abs(text.hashCode()))}"
+                    return materialize(suggestedName = "$prefix$name")
+                }
+
+                when (it) {
+                    is SimpleColumnGroup -> {
+                        val nestedSchema = PluginDataFrameSchema(it.columns()).materialize(it)
+                        val columnsContainerReturnType =
+                            ConeClassLikeTypeImpl(
+                                ConeClassLikeLookupTagImpl(Names.COLUM_GROUP_CLASS_ID),
+                                typeArguments = arrayOf(nestedSchema.schema.defaultType()),
+                                isNullable = false
+                            )
+
+                        val dataRowReturnType =
+                            ConeClassLikeTypeImpl(
+                                ConeClassLikeLookupTagImpl(Names.DATA_ROW_CLASS_ID),
+                                typeArguments = arrayOf(nestedSchema.schema.defaultType()),
+                                isNullable = false
+                            )
+
+                        SchemaProperty(schema.defaultType(), PropertyName.of(it.name), dataRowReturnType, columnsContainerReturnType)
+                    }
+
+                    is SimpleFrameColumn -> {
+                        val nestedClassMarker = PluginDataFrameSchema(it.columns()).materialize(it)
+                        val frameColumnReturnType =
+                            ConeClassLikeTypeImpl(
+                                ConeClassLikeLookupTagImpl(Names.DF_CLASS_ID),
+                                typeArguments = arrayOf(nestedClassMarker.schema.defaultType()),
+                                isNullable = false
+                            )
+
+                        SchemaProperty(
+                            marker = schema.defaultType(),
+                            propertyName = PropertyName.of(it.name),
+                            dataRowReturnType = frameColumnReturnType,
+                            columnContainerReturnType = frameColumnReturnType.toFirResolvedTypeRef()
+                                .projectOverDataColumnType()
+                        )
+                    }
+
+                    is SimpleDataColumn -> SchemaProperty(
+                        marker = schema.defaultType(),
+                        propertyName = PropertyName.of(it.name),
+                        dataRowReturnType = it.type.type(),
+                        columnContainerReturnType = it.type.type().toFirResolvedTypeRef().projectOverDataColumnType()
+                    )
+                }
+            }
+            schema.callShapeData = CallShapeData.Schema(properties)
+            scope.callShapeData = CallShapeData.Scope(properties, call.calleeReference.source)
+            val schemaApi = DataSchemaApi(schema, scope)
+            dataSchemaApis.add(schemaApi)
+            return schemaApi
+        }
+
+        dataFrameSchema.materialize(firstSchema)
+        return dataSchemaApis
+    }
+
+    data class DataSchemaApi(val schema: FirRegularClass, val scope: FirRegularClass)
+
+    private fun buildSchema(tokenId: ClassId): FirRegularClass {
+        val token = buildRegularClass {
+            moduleData = session.moduleData
+            resolvePhase = FirResolvePhase.BODY_RESOLVE
+            origin = FirDeclarationOrigin.Source
+            status = FirResolvedDeclarationStatusImpl(Visibilities.Local, Modality.ABSTRACT, EffectiveVisibility.Local)
+            deprecationsProvider = EmptyDeprecationsProvider
+            classKind = ClassKind.CLASS
+            scopeProvider = FirKotlinScopeProvider()
+            superTypeRefs += FirImplicitAnyTypeRef(null)
+
+            name = tokenId.shortClassName
+            this.symbol = FirRegularClassSymbol(tokenId)
+        }
+        return token
+    }
+
+    private fun findLet(): FirFunctionSymbol<*> {
+        return session.symbolProvider.getTopLevelFunctionSymbols(FqName("kotlin"), Name.identifier("let")).single()
+    }
+
+    private fun findRun(): FirFunctionSymbol<*> {
+        return session.symbolProvider.getTopLevelFunctionSymbols(FqName("kotlin"), Name.identifier("run")).single { it.typeParameterSymbols.size == 1 }
+    }
+
+    private fun String.titleCase() = replaceFirstChar { it.uppercaseChar() }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/IrBodyFiller.kt

@@ -0,0 +1,211 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlinx.dataframe.plugin.extensions
+
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.backend.js.utils.valueArguments
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrConstructor
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationWithName
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.declarations.IrPackageFragment
+import org.jetbrains.kotlin.ir.declarations.IrProperty
+import org.jetbrains.kotlin.ir.declarations.copyAttributes
+import org.jetbrains.kotlin.ir.declarations.createBlockBody
+import org.jetbrains.kotlin.ir.expressions.IrBody
+import org.jetbrains.kotlin.ir.expressions.IrConst
+import org.jetbrains.kotlin.ir.expressions.IrErrorCallExpression
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
+import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
+import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrDelegatingConstructorCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrTypeOperatorCallImpl
+import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
+import org.jetbrains.kotlin.ir.symbols.UnsafeDuringIrConstructionAPI
+import org.jetbrains.kotlin.ir.types.IrSimpleType
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.classFqName
+import org.jetbrains.kotlin.ir.types.classOrFail
+import org.jetbrains.kotlin.ir.types.classifierOrNull
+import org.jetbrains.kotlin.ir.types.getClass
+import org.jetbrains.kotlin.ir.util.constructors
+import org.jetbrains.kotlin.ir.util.findAnnotation
+import org.jetbrains.kotlin.ir.util.parentAsClass
+import org.jetbrains.kotlin.ir.util.primaryConstructor
+import org.jetbrains.kotlin.ir.util.superTypes
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+
+class IrBodyFiller : IrGenerationExtension {
+    override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
+        DataFrameFileLowering(pluginContext).lower(moduleFragment)
+    }
+}
+
+private class DataFrameFileLowering(val context: IrPluginContext) : FileLoweringPass, IrElementTransformerVoid() {
+    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"))
+    }
+
+    override fun lower(irFile: IrFile) {
+        irFile.transformChildren(this, null)
+    }
+
+    override fun visitConstructor(declaration: IrConstructor): IrStatement {
+        val origin = declaration.origin
+        if (!(origin is IrDeclarationOrigin.GeneratedByPlugin && origin.pluginKey is TokenGenerator.Key)) return declaration
+        declaration.body = generateBodyForDefaultConstructor(declaration)
+        return declaration
+    }
+
+    @OptIn(UnsafeDuringIrConstructionAPI::class)
+    private fun generateBodyForDefaultConstructor(declaration: IrConstructor): IrBody? {
+        val irType = declaration.returnType.superTypes()[0]
+        val symbol = irType.classOrFail.owner.primaryConstructor?.symbol ?: return null
+        val type = declaration.returnType as? IrSimpleType ?: return null
+        val delegatingAnyCall = IrDelegatingConstructorCallImpl(
+            -1,
+            -1,
+            irType,
+            symbol,
+            typeArgumentsCount = 0,
+            valueArgumentsCount = 0
+        ).copyAttributes(declaration.parentAsClass)
+
+        val initializerCall = IrInstanceInitializerCallImpl(
+            -1,
+            -1,
+            (declaration.parent as? IrClass)?.symbol ?: return null,
+            type
+        )
+
+        return context.irFactory.createBlockBody(-1, -1, listOf(delegatingAnyCall, initializerCall))
+    }
+
+    @OptIn(UnsafeDuringIrConstructionAPI::class)
+    override fun visitProperty(declaration: IrProperty): IrStatement {
+        val origin = declaration.origin
+        val pluginKey = (origin as? IrDeclarationOrigin.GeneratedByPlugin)?.pluginKey as? DataFramePlugin
+        if (pluginKey == null) {
+            declaration.transformChildren(this, null)
+            return declaration
+        }
+        val getter = declaration.getter ?: return declaration
+
+        val constructors = context.referenceConstructors(ClassId(FqName("kotlin.jvm"), Name.identifier("JvmName")))
+        val jvmName = constructors.single { it.owner.valueParameters.size == 1 }
+        val marker =
+            ((getter.extensionReceiverParameter!!.type as IrSimpleType).arguments.single() as IrSimpleType).classOrFail.owner
+        val jvmNameArg = "${marker.nestedName()}_${declaration.name.identifier}"
+        getter.annotations = listOf(
+            IrConstructorCallImpl(-1, -1, jvmName.owner.returnType, jvmName, 0, 0, 1)
+                .also {
+                    it.putValueArgument(0, IrConstImpl.string(-1, -1, context.irBuiltIns.stringType, jvmNameArg))
+                }
+        )
+        val returnType = getter.returnType
+        val isDataColumn = returnType.classFqName!!.asString().let {
+            it == DataColumn::class.qualifiedName!! || it == ColumnGroup::class.qualifiedName!!
+        }
+
+        val get = if (isDataColumn) {
+            context
+                .referenceFunctions(COLUMNS_SCOPE_ID)
+                .single {
+                    it.owner.valueParameters.size == 1 && it.owner.valueParameters[0].type == context.irBuiltIns.stringType
+                }
+        } else {
+            context
+                .referenceFunctions(DATA_ROW_ID)
+                .single {
+                    it.owner.valueParameters.size == 1 && it.owner.valueParameters[0].type == context.irBuiltIns.stringType
+                }
+        }
+
+        val call = IrCallImpl(-1, -1, context.irBuiltIns.anyNType, get, 0, 1).also {
+            val thisSymbol: IrValueSymbol = getter.extensionReceiverParameter?.symbol!!
+            it.dispatchReceiver = IrGetValueImpl(-1, -1, thisSymbol)
+            val annotation = declaration.annotations.findAnnotation(Names.COLUMN_NAME_ANNOTATION.asSingleFqName())
+            val columnName = (annotation?.valueArguments?.get(0) as? IrConst<*>)?.value as? String
+            val columName = columnName ?: declaration.name.identifier
+            it.putValueArgument(0, IrConstImpl.string(-1, -1, context.irBuiltIns.stringType, columName))
+        }
+
+        val typeOp = IrTypeOperatorCallImpl(-1, -1, returnType, IrTypeOperator.CAST, returnType, call)
+        val returnExpression = IrReturnImpl(-1, -1, returnType, getter.symbol, typeOp)
+        getter.apply {
+            body = factory.createBlockBody(-1, -1, listOf(returnExpression))
+        }
+
+        return declaration
+    }
+
+    private fun IrDeclarationWithName.nestedName() = buildString { computeNestedName(this@nestedName, this) }
+
+    private fun computeNestedName(declaration: IrDeclarationWithName, result: StringBuilder): Boolean {
+        when (val parent = declaration.parent) {
+            is IrClass -> {
+                if (!computeNestedName(parent, result)) return false
+            }
+            is IrPackageFragment -> {}
+            else -> return false
+        }
+        if (result.isNotEmpty()) result.append('_')
+        result.append(declaration.name.asString())
+        return true
+    }
+
+    // org.jetbrains.kotlin.fir.backend.generators.CallAndReferenceGenerator#applyReceivers
+    override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression {
+        if (isScope(expression.typeOperand)) {
+            return expression.replaceWithConstructorCall()
+        }
+        return super.visitTypeOperator(expression)
+    }
+
+    override fun visitErrorCallExpression(expression: IrErrorCallExpression): IrExpression {
+        if (!isScope(expression.type)) {
+            return expression
+        }
+        return expression.replaceWithConstructorCall()
+    }
+
+    @OptIn(UnsafeDuringIrConstructionAPI::class)
+    private fun isScope(type: IrType): Boolean {
+        val origin = (type.classifierOrNull?.owner as? IrClass)?.origin ?: return false
+        val fromPlugin = origin is IrDeclarationOrigin.GeneratedByPlugin && origin.pluginKey is DataFramePlugin
+        val scopeReference = type.classFqName?.shortName()?.asString()?.startsWith("Scope") ?: false
+        return fromPlugin || scopeReference
+    }
+
+    @OptIn(UnsafeDuringIrConstructionAPI::class)
+    private fun IrExpression.replaceWithConstructorCall(): IrConstructorCallImpl {
+        val constructor = type.getClass()!!.constructors.toList().single()
+        return IrConstructorCallImpl(-1, -1, type, constructor.symbol, 0, 0, 0)
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/KotlinTypeFacade.kt

@@ -0,0 +1,96 @@
+package org.jetbrains.kotlinx.dataframe.plugin.extensions
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.ConeFlexibleType
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeNullability
+import org.jetbrains.kotlin.fir.types.isNullable
+import org.jetbrains.kotlin.fir.types.typeContext
+import org.jetbrains.kotlin.fir.types.withNullability
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.name.StandardClassIds
+
+interface KotlinTypeFacade : SessionContext {
+    val isTest: Boolean
+
+    fun Marker.type() = type
+
+    fun Marker.changeNullability(map: (Boolean) -> Boolean): Marker {
+        val coneNullability = when (map(type.isNullable)) {
+            true -> ConeNullability.NULLABLE
+            false -> ConeNullability.NOT_NULL
+        }
+        return Marker(type = type.withNullability(coneNullability, session.typeContext))
+    }
+
+    fun Marker.isList(): Boolean {
+        return type.isBuiltinType(List, isNullable = null)
+    }
+
+    fun Marker.typeArgument(): Marker {
+        val argument = when (val argument = type.typeArguments[0]) {
+            is ConeKotlinType -> argument
+            else -> error("${argument::class} ${argument}")
+        }
+        return Marker(argument)
+    }
+}
+
+interface SessionContext {
+    val session: FirSession
+}
+
+fun SessionContext(session: FirSession) = object : SessionContext {
+    override val session: FirSession = session
+}
+
+private val List = "List".collectionsId()
+
+private fun ConeKotlinType.isBuiltinType(classId: ClassId, isNullable: Boolean?): Boolean {
+    if (this !is ConeClassLikeType) return false
+    return lookupTag.classId == classId && (isNullable == null || type.isNullable == isNullable)
+}
+
+private fun String.collectionsId() = ClassId(StandardClassIds.BASE_COLLECTIONS_PACKAGE, Name.identifier(this))
+
+class KotlinTypeFacadeImpl(
+    override val session: FirSession,
+    override val isTest: Boolean
+) : KotlinTypeFacade
+
+class Marker private constructor(internal val type: ConeKotlinType) {
+    companion object {
+        operator fun invoke(type: ConeKotlinType): Marker {
+            val type = if (type is ConeFlexibleType) {
+                type.lowerBound
+            } else {
+                type
+            }
+            return Marker(type)
+        }
+    }
+
+    override fun toString(): String {
+        return "Marker(type=$type (${type::class}))"
+    }
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as Marker
+
+        return type == other.type
+    }
+
+    override fun hashCode(): Int {
+        return type.hashCode()
+    }
+}
+
+fun ConeKotlinType.wrap(): Marker = Marker(this)
+
+
+

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ReturnTypeBasedReceiverInjector.kt

@@ -0,0 +1,50 @@
+package org.jetbrains.kotlinx.dataframe.plugin.extensions
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.extensions.FirExpressionResolutionExtension
+import org.jetbrains.kotlin.fir.scopes.collectAllProperties
+import org.jetbrains.kotlin.fir.scopes.impl.declaredMemberScope
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.resolvedType
+import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+
+class ReturnTypeBasedReceiverInjector(session: FirSession) : FirExpressionResolutionExtension(session) {
+    companion object {
+        private val SCHEMA_TYPES = setOf(
+            Names.DF_CLASS_ID,
+            Names.GROUP_BY_CLASS_ID,
+            Names.DATA_ROW_CLASS_ID,
+            Names.COLUM_GROUP_CLASS_ID,
+        )
+    }
+
+    @OptIn(SymbolInternals::class)
+    override fun addNewImplicitReceivers(functionCall: FirFunctionCall): List<ConeKotlinType> {
+        val callReturnType = functionCall.resolvedType
+        return if (callReturnType.classId in SCHEMA_TYPES) {
+            val typeArguments = callReturnType.typeArguments
+            typeArguments
+                .mapNotNull {
+                    val symbol = (it as? ConeKotlinType)?.toRegularClassSymbol(session)
+                    symbol?.takeIf { it.fir.callShapeData != null }
+                }
+                .takeIf { it.size == typeArguments.size }
+                .orEmpty()
+                .flatMap { marker ->
+                    marker.declaredMemberScope(session, FirResolvePhase.DECLARATIONS).collectAllProperties()
+                        .filterIsInstance<FirPropertySymbol>()
+                        .filter { it.getAnnotationByClassId(Names.SCOPE_PROPERTY_ANNOTATION, session) != null }
+                        .map { it.resolvedReturnType }
+                }
+        } else {
+            emptyList()
+        }
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TokenGenerator.kt

@@ -0,0 +1,151 @@
+package org.jetbrains.kotlinx.dataframe.plugin.extensions
+
+import org.jetbrains.kotlin.GeneratedDeclarationKey
+import org.jetbrains.kotlin.descriptors.EffectiveVisibility
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.caches.FirCache
+import org.jetbrains.kotlin.fir.caches.createCache
+import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.caches.getValue
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+import org.jetbrains.kotlinx.dataframe.plugin.utils.generateExtensionProperty
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
+import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
+import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationArgumentMapping
+import org.jetbrains.kotlin.fir.expressions.builder.buildLiteralExpression
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
+import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
+import org.jetbrains.kotlin.fir.plugin.createConstructor
+import org.jetbrains.kotlin.fir.plugin.createMemberProperty
+import org.jetbrains.kotlin.fir.resolve.defaultType
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.toFirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.name.SpecialNames
+import org.jetbrains.kotlin.types.ConstantValueKind
+
+class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
+    object Key : GeneratedDeclarationKey()
+
+    @OptIn(SymbolInternals::class)
+    private val propertiesCache: FirCache<FirClassSymbol<*>, Map<Name, List<FirProperty>>?, Nothing?> =
+        session.firCachesFactory.createCache { k ->
+            val callShapeData = k.fir.callShapeData ?: return@createCache null
+            when (callShapeData) {
+                is CallShapeData.Schema -> callShapeData.columns.withIndex().associate { (index, property) ->
+                    val resolvedTypeRef = buildResolvedTypeRef {
+                        type = property.dataRowReturnType
+                    }
+                    val identifier = property.propertyName.identifier
+                    identifier to listOf(buildProperty(resolvedTypeRef, identifier, k, property.propertyName.columnNameAnnotation, order = index))
+                }
+                is CallShapeData.RefinedType -> callShapeData.scopes.associate {
+                    val identifier = Name.identifier(it.name.identifier.replaceFirstChar { it.lowercaseChar() })
+                    identifier to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), identifier, k, isScopeProperty = true))
+                }
+                is CallShapeData.Scope -> callShapeData.columns.associate { schemaProperty ->
+                    val propertyName = schemaProperty.propertyName
+                    val callableId = CallableId(k.classId, propertyName.identifier)
+                    val dataRowExtension = generateExtensionProperty(
+                        callableId = callableId,
+                        receiverType = ConeClassLikeTypeImpl(
+                            ConeClassLikeLookupTagImpl(Names.DATA_ROW_CLASS_ID),
+                            typeArguments = arrayOf(schemaProperty.marker),
+                            isNullable = false
+                        ),
+                        propertyName = propertyName,
+                        returnTypeRef = schemaProperty.dataRowReturnType.toFirResolvedTypeRef(),
+                        symbol = k,
+                        effectiveVisibility = EffectiveVisibility.Local,
+                        source = callShapeData.source
+                    )
+
+                    val columnContainerExtension = generateExtensionProperty(
+                        callableId = callableId,
+                        receiverType = ConeClassLikeTypeImpl(
+                            ConeClassLikeLookupTagImpl(Names.COLUMNS_SCOPE_CLASS_ID),
+                            typeArguments = arrayOf(schemaProperty.marker),
+                            isNullable = false
+                        ),
+                        propertyName = propertyName,
+                        returnTypeRef = schemaProperty.columnContainerReturnType.toFirResolvedTypeRef(),
+                        symbol = k,
+                        effectiveVisibility = EffectiveVisibility.Local,
+                        source = callShapeData.source
+                    )
+                    propertyName.identifier to listOf(dataRowExtension, columnContainerExtension)
+                }
+            }
+        }
+
+    @OptIn(SymbolInternals::class)
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
+        val destination = mutableSetOf<Name>()
+        when (classSymbol.fir.callShapeData) {
+            is CallShapeData.RefinedType -> destination.add(SpecialNames.INIT)
+            is CallShapeData.Schema -> destination.add(SpecialNames.INIT)
+            is CallShapeData.Scope -> destination.add(SpecialNames.INIT)
+            null -> Unit
+        }
+        return propertiesCache.getValue(classSymbol)?.values?.flatten()?.mapTo(destination) { it.name } ?: emptySet()
+    }
+
+    override fun generateProperties(callableId: CallableId, context: MemberGenerationContext?): List<FirPropertySymbol> {
+        val owner = context?.owner ?: return emptyList()
+        val properties = propertiesCache.getValue(owner)?.get(callableId.callableName) ?: return emptyList()
+        return properties.map { it.symbol }
+    }
+
+    private fun buildProperty(
+        resolvedTypeRef: FirResolvedTypeRef,
+        propertyName: Name,
+        k: FirClassSymbol<*>,
+        columnNameAnnotation: FirAnnotation? = null,
+        isScopeProperty: Boolean = false,
+        order: Int? = null,
+    ): FirProperty {
+        return createMemberProperty(k, Key, propertyName, resolvedTypeRef.type) {
+            modality = Modality.ABSTRACT
+            visibility = Visibilities.Public
+        }.apply {
+            val annotations = mutableListOf<FirAnnotation>()
+            if (order != null) {
+                annotations += buildAnnotation {
+                    annotationTypeRef = buildResolvedTypeRef {
+                        type = Names.ORDER_ANNOTATION.defaultType(emptyList())
+                    }
+                    argumentMapping = buildAnnotationArgumentMapping {
+                        mapping[Names.ORDER_ARGUMENT] = buildLiteralExpression(null, ConstantValueKind.Int, order, setType = true)
+                    }
+                }
+            }
+            if (isScopeProperty) {
+                annotations += buildAnnotation {
+                    annotationTypeRef = buildResolvedTypeRef {
+                        type = Names.SCOPE_PROPERTY_ANNOTATION.defaultType(emptyList())
+                    }
+                    argumentMapping = buildAnnotationArgumentMapping()
+                }
+            }
+            columnNameAnnotation?.let {
+                annotations += it
+            }
+            replaceAnnotations(annotations)
+        }
+    }
+
+    override fun generateConstructors(context: MemberGenerationContext): List<FirConstructorSymbol> {
+        return listOf(createConstructor(context.owner, Key, isPrimary = true).symbol)
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TopLevelExtensionsGenerator.kt

@@ -0,0 +1,155 @@
+package org.jetbrains.kotlinx.dataframe.plugin.extensions
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.checkers.toClassLikeSymbol
+import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+import org.jetbrains.kotlinx.dataframe.plugin.utils.generateExtensionProperty
+import org.jetbrains.kotlinx.dataframe.plugin.utils.projectOverDataColumnType
+import org.jetbrains.kotlin.fir.declarations.hasAnnotation
+import org.jetbrains.kotlin.fir.declarations.utils.isLocal
+import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
+import org.jetbrains.kotlin.fir.extensions.ExperimentalTopLevelDeclarationsGenerationApi
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
+import org.jetbrains.kotlin.fir.extensions.predicate.LookupPredicate
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
+import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.toFirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.ConeTypeProjection
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.constructType
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.fir.types.toSymbol
+import org.jetbrains.kotlin.fir.types.toTypeProjection
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.PropertyName
+
+/**
+ * extensions inside scope classes are generated here:
+ * @see org.jetbrains.kotlinx.dataframe.plugin.extensions.TokenGenerator
+ */
+class TopLevelExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
+    private companion object {
+        val dataSchema = FqName(DataSchema::class.qualifiedName!!)
+    }
+
+    private val predicateBasedProvider = session.predicateBasedProvider
+    private val matchedClasses by lazy {
+        predicateBasedProvider.getSymbolsByPredicate(predicate).filterIsInstance<FirRegularClassSymbol>()
+    }
+
+    private val predicate: LookupPredicate = LookupPredicate.BuilderContext.annotated(dataSchema)
+
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(predicate)
+    }
+
+    private val fields by lazy {
+        matchedClasses.filterNot { it.isLocal }.flatMap { classSymbol ->
+            classSymbol.declarationSymbols.filterIsInstance<FirPropertySymbol>().map { propertySymbol ->
+                val callableId = propertySymbol.callableId
+                DataSchemaField(
+                    classSymbol,
+                    propertySymbol,
+                    CallableId(packageName = callableId.packageName, className = null, callableName = callableId.callableName)
+                )
+            }
+        }
+    }
+
+    private data class DataSchemaField(
+        val classSymbol: FirRegularClassSymbol,
+        val propertySymbol: FirPropertySymbol,
+        val callableId: CallableId
+    )
+
+    @OptIn(ExperimentalTopLevelDeclarationsGenerationApi::class)
+    override fun getTopLevelCallableIds(): Set<CallableId> {
+        return buildSet {
+            fields.mapTo(this) { it.callableId }
+        }
+    }
+
+    override fun generateProperties(callableId: CallableId, context: MemberGenerationContext?): List<FirPropertySymbol> {
+        val owner = context?.owner
+        return when (owner) {
+            null -> fields.filter { it.callableId == callableId }.flatMap { (owner, property, callableId) ->
+                var resolvedReturnTypeRef = property.resolvedReturnTypeRef
+                val columnName = property.getAnnotationByClassId(Names.COLUMN_NAME_ANNOTATION, session)?.let { annotation ->
+                    (annotation.argumentMapping.mapping[Names.COLUMN_NAME_ARGUMENT] as? FirLiteralExpression)?.value as? String?
+                }
+                val name = property.name
+                val marker = owner.constructType(arrayOf(), isNullable = false).toTypeProjection(Variance.INVARIANT)
+
+                val columnGroupProjection: ConeTypeProjection? = if (resolvedReturnTypeRef.coneType.classId?.equals(
+                        Names.DATA_ROW_CLASS_ID) == true) {
+                    resolvedReturnTypeRef.coneType.typeArguments[0]
+                } else if (resolvedReturnTypeRef.toClassLikeSymbol(session)?.hasAnnotation(Names.DATA_SCHEMA_CLASS_ID, session) == true) {
+                    resolvedReturnTypeRef.coneType
+                } else {
+                    null
+                }
+
+                if (
+                    resolvedReturnTypeRef.type.classId?.equals(Names.LIST) == true &&
+                    (resolvedReturnTypeRef.type.typeArguments[0] as? ConeClassLikeType)?.toSymbol(session)?.hasAnnotation(
+                        Names.DATA_SCHEMA_CLASS_ID, session) == true
+                ) {
+                    require(columnGroupProjection == null)
+                    resolvedReturnTypeRef = ConeClassLikeTypeImpl(
+                        ConeClassLikeLookupTagImpl(Names.DF_CLASS_ID),
+                        typeArguments = arrayOf(resolvedReturnTypeRef.type.typeArguments[0]),
+                        isNullable = false
+                    ).toFirResolvedTypeRef()
+                }
+
+                val rowExtension = generateExtensionProperty(
+                    callableId = callableId,
+                    receiverType = ConeClassLikeTypeImpl(
+                        ConeClassLikeLookupTagImpl(Names.DATA_ROW_CLASS_ID),
+                        typeArguments = arrayOf(marker),
+                        isNullable = false
+                    ),
+                    propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }),
+                    returnTypeRef = resolvedReturnTypeRef,
+                    source = owner.source
+                )
+
+                val columnReturnType = when {
+                    columnGroupProjection != null -> {
+                        ConeClassLikeTypeImpl(
+                            ConeClassLikeLookupTagImpl(Names.COLUM_GROUP_CLASS_ID),
+                            typeArguments = arrayOf(columnGroupProjection),
+                            isNullable = false
+                        ).toFirResolvedTypeRef()
+                    }
+
+                    else -> resolvedReturnTypeRef.projectOverDataColumnType().toFirResolvedTypeRef()
+                }
+                val columnsContainerExtension = generateExtensionProperty(
+                    callableId = callableId,
+                    receiverType = ConeClassLikeTypeImpl(
+                        ConeClassLikeLookupTagImpl(Names.COLUMNS_CONTAINER_CLASS_ID),
+                        typeArguments = arrayOf(marker),
+                        isNullable = false
+                    ),
+                    propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }),
+                    returnTypeRef = columnReturnType,
+                    source = owner.source
+                )
+                listOf(rowExtension.symbol, columnsContainerExtension.symbol)
+            }
+
+            else -> emptyList()
+        }
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/impl/PropertyName.kt

@@ -0,0 +1,52 @@
+package org.jetbrains.kotlinx.dataframe.plugin.extensions.impl
+
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
+import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationArgumentMapping
+import org.jetbrains.kotlin.fir.expressions.builder.buildLiteralExpression
+import org.jetbrains.kotlin.fir.resolve.defaultType
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.types.ConstantValueKind
+import org.jetbrains.kotlinx.dataframe.codeGen.ValidFieldName
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+
+data class PropertyName(val identifier: Name, val columnNameAnnotation: FirAnnotation?) {
+    companion object {
+        fun of(name: String): PropertyName {
+            val valid = ValidFieldName.of(name)
+            var columnName = false
+            val identifier = if (valid.unquoted != name) {
+                columnName = true
+                Name.identifier(valid.unquoted)
+            } else {
+                Name.identifier(name)
+            }
+            val columnNameAnnotation: FirAnnotation? = if (columnName) {
+                buildAnnotation(name)
+            } else {
+                null
+            }
+            return PropertyName(identifier, columnNameAnnotation)
+        }
+
+        fun buildAnnotation(name: String): FirAnnotation {
+            return org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation {
+                annotationTypeRef = buildResolvedTypeRef {
+                    type = Names.COLUMN_NAME_ANNOTATION.defaultType(emptyList())
+                }
+                argumentMapping = buildAnnotationArgumentMapping {
+                    mapping[Names.COLUMN_NAME_ARGUMENT] = buildLiteralExpression(
+                        source = null,
+                        kind = ConstantValueKind.String,
+                        value = name,
+                        setType = true
+                    )
+                }
+            }
+        }
+
+        fun of(identifier: Name, columnNameAnnotation: FirAnnotation?): PropertyName {
+            return PropertyName(identifier, columnNameAnnotation)
+        }
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/impl/SchemaProperty.kt

@@ -0,0 +1,12 @@
+package org.jetbrains.kotlinx.dataframe.plugin.extensions.impl
+
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeTypeProjection
+
+data class SchemaProperty(
+    val marker: ConeTypeProjection,
+    val propertyName: PropertyName,
+    val dataRowReturnType: ConeKotlinType,
+    val columnContainerReturnType: ConeKotlinType,
+    val override: Boolean = false
+)

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/DataFrameAdapter.kt

@@ -0,0 +1,54 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl
+
+import org.jetbrains.kotlinx.dataframe.AnyCol
+import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.api.asDataColumn
+import org.jetbrains.kotlinx.dataframe.api.cast
+import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
+import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
+import org.jetbrains.kotlinx.dataframe.columns.FrameColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
+
+fun PluginDataFrameSchema.asDataFrame(): DataFrame<ConeTypesAdapter> {
+    val df = columns().map()
+    return df
+}
+
+fun DataFrame<ConeTypesAdapter>.toPluginDataFrameSchema() = PluginDataFrameSchema(columns().mapBack())
+
+interface ConeTypesAdapter
+
+private fun List<SimpleCol>.map(): DataFrame<ConeTypesAdapter> {
+    val columns = map {
+        it.asDataColumn()
+    }
+    return dataFrameOf(columns).cast()
+}
+
+fun SimpleCol.asDataColumn(): DataColumn<*> {
+    val column = when (this) {
+        is SimpleDataColumn -> DataColumn.createByType(this.name, listOf(this.type))
+        is SimpleColumnGroup -> DataColumn.createColumnGroup(this.name, this.columns().map()).asDataColumn()
+        is SimpleFrameColumn -> DataColumn.createFrameColumn(this.name, listOf(this.columns().map()))
+    }
+    return column
+}
+
+private fun List<AnyCol>.mapBack(): List<SimpleCol> = map { it.asSimpleColumn() }
+
+fun AnyCol.asSimpleColumn(): SimpleCol {
+    return when (this) {
+        is ColumnGroup<*> -> {
+            SimpleColumnGroup(name(), columns().mapBack())
+        }
+
+        is FrameColumn<*> -> {
+            SimpleFrameColumn(name(), this[0].columns().mapBack())
+        }
+
+        else -> {
+            SimpleDataColumn(name(), this[0] as TypeApproximation)
+        }
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/ExpectedArgumentDelegates.kt

@@ -0,0 +1,42 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter.*
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBy
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.DataFrameCallableId
+import kotlin.properties.PropertyDelegateProvider
+import kotlin.properties.ReadOnlyProperty
+
+typealias ExpectedArgumentProvider<T> = PropertyDelegateProvider<Any?, ReadOnlyProperty<Arguments, T>>
+
+fun <T> AbstractInterpreter<T>.dataFrame(
+    name: ArgumentName? = null
+): ExpectedArgumentProvider<PluginDataFrameSchema> = arg(name, lens = Interpreter.Schema)
+
+fun <T> AbstractInterpreter<T>.type(
+    name: ArgumentName? = null
+): ExpectedArgumentProvider<TypeApproximation> = arg(name, lens = Interpreter.ReturnType)
+
+fun <T, E : Enum<E>> AbstractInterpreter<T>.enum(
+    name: ArgumentName? = null,
+    defaultValue: DefaultValue<E> = Absent
+): ExpectedArgumentProvider<E> = argConvert(name = name, defaultValue = defaultValue) { it: DataFrameCallableId ->
+    val forName: Class<*> = Class.forName("${it.packageName}.${it.className}")
+    @Suppress("UNCHECKED_CAST")
+    java.lang.Enum.valueOf(forName as Class<out Enum<*>>, it.callableName) as E
+}
+
+internal fun <T> AbstractInterpreter<T>.dsl(
+    name: ArgumentName? = null
+): ExpectedArgumentProvider<(Any, Map<String, Interpreter.Success<Any?>>) -> Unit> =
+    arg(name, lens = Interpreter.Dsl, defaultValue = Present(value = {_, _ -> }))
+
+internal fun <T> AbstractInterpreter<T>.ignore(
+    name: ArgumentName? = null
+): ExpectedArgumentProvider<Nothing?> =
+    arg(name, lens = Interpreter.Id, defaultValue = Present(null))
+
+internal fun <T> AbstractInterpreter<T>.groupBy(
+    name: ArgumentName? = null
+): ExpectedArgumentProvider<GroupBy> = arg(name, lens = Interpreter.GroupBy)
+

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/Interpreter.kt

@@ -0,0 +1,150 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl
+
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
+import kotlin.properties.PropertyDelegateProvider
+import kotlin.properties.ReadOnlyProperty
+import kotlin.reflect.KType
+import kotlin.reflect.typeOf
+
+interface Interpreter<T> {
+    val expectedArguments: List<ExpectedArgument>
+
+    data class ExpectedArgument(
+        val name: String,
+        val klass: KType,
+        val lens: Lens,
+        val defaultValue: DefaultValue<*>
+    )
+
+    sealed interface Lens
+
+    data object Value : Lens
+
+    data object ReturnType : Lens
+
+    data object Dsl : Lens
+
+    data object Schema : Lens
+
+    data object GroupBy : Lens
+
+    data object  Id : Lens
+
+    // required to compute whether resulting schema should be inheritor of previous class or a new class
+    fun startingSchema(arguments: Map<String, Success<Any?>>, kotlinTypeFacade: KotlinTypeFacade): PluginDataFrameSchema?
+
+    fun interpret(arguments: Map<String, Success<Any?>>, kotlinTypeFacade: KotlinTypeFacade): InterpretationResult<T>
+
+    sealed interface InterpretationResult<out T>
+
+    class Success<out T>(val value: T) : InterpretationResult<T> {
+        override fun equals(other: Any?): Boolean {
+            if (this === other) return true
+            if (javaClass != other?.javaClass) return false
+
+            other as Success<*>
+
+            return value == other.value
+        }
+
+        override fun hashCode(): Int {
+            return value?.hashCode() ?: 0
+        }
+    }
+
+    class Error(val message: String?) : InterpretationResult<Nothing>
+}
+
+sealed interface DefaultValue<out T>
+
+class Present<T>(val value: T) : DefaultValue<T>
+data object Absent : DefaultValue<Nothing>
+
+open class Arguments(private val arguments: Map<String, Interpreter.Success<Any?>>, kotlinTypeFacade: KotlinTypeFacade): KotlinTypeFacade by kotlinTypeFacade {
+    operator fun get(s: String): Any? = (arguments[s] ?: error("")).value
+    operator fun contains(key: String): Boolean {
+        return arguments.contains(key)
+    }
+}
+
+abstract class AbstractInterpreter<T> : Interpreter<T> {
+    @PublishedApi
+    internal val _expectedArguments: MutableList<Interpreter.ExpectedArgument> = mutableListOf()
+
+    override val expectedArguments: List<Interpreter.ExpectedArgument> = _expectedArguments
+
+    protected open val Arguments.startingSchema: PluginDataFrameSchema? get() = null
+
+    final override fun startingSchema(arguments: Map<String, Interpreter.Success<Any?>>, kotlinTypeFacade: KotlinTypeFacade): PluginDataFrameSchema? {
+        return Arguments(arguments, kotlinTypeFacade).startingSchema
+    }
+
+    inline fun <Value, reified CompileTimeValue> argConvert(
+        defaultValue: DefaultValue<Value> = Absent,
+        name: ArgumentName? = null,
+        lens: Interpreter.Lens = Interpreter.Value,
+        crossinline converter: (CompileTimeValue) -> Value
+    ): PropertyDelegateProvider<Any?, ReadOnlyProperty<Arguments, Value>> = PropertyDelegateProvider { thisRef: Any?, property ->
+        val name = name?.value ?: property.name
+        _expectedArguments.add(Interpreter.ExpectedArgument(name, typeOf<CompileTimeValue>(), lens, defaultValue))
+        ReadOnlyProperty { args, _ ->
+            if (name !in args && defaultValue is Present) {
+                defaultValue.value
+            } else {
+                converter(args[name] as CompileTimeValue)
+            }
+        }
+    }
+
+    fun <Value> arg(
+        name: ArgumentName? = null,
+        expectedType: KType? = null,
+        defaultValue: DefaultValue<Value> = Absent,
+        lens: Interpreter.Lens = Interpreter.Value
+    ): PropertyDelegateProvider<Any?, ReadOnlyProperty<Arguments, Value>> = PropertyDelegateProvider { thisRef: Any?, property ->
+        val name = name?.value ?: property.name
+        _expectedArguments.add(
+            Interpreter.ExpectedArgument(
+                name,
+                expectedType ?: property.returnType,
+                lens,
+                defaultValue
+            )
+        )
+        ReadOnlyProperty { args, _ ->
+            if (name !in args && defaultValue is Present) {
+                defaultValue.value
+            } else {
+                @Suppress("UNCHECKED_CAST")
+                args[name] as Value
+            }
+        }
+    }
+
+    class ArgumentName private constructor(val value: String) {
+        companion object {
+            fun of(name: String): ArgumentName = ArgumentName(name)
+        }
+    }
+
+    fun name(name: String): ArgumentName = ArgumentName.of(name)
+
+    final override fun interpret(arguments: Map<String, Interpreter.Success<Any?>>, kotlinTypeFacade: KotlinTypeFacade): Interpreter.InterpretationResult<T> {
+        return try {
+            Arguments(arguments, kotlinTypeFacade).interpret().let { Interpreter.Success(it) }
+        } catch (e: Exception) {
+            Interpreter.Error(e.message + e.stackTrace.contentToString())
+        }
+    }
+
+    abstract fun Arguments.interpret(): T
+}
+
+interface SchemaModificationInterpreter : Interpreter<PluginDataFrameSchema> {
+
+    override fun interpret(arguments: Map<String, Interpreter.Success<Any?>>, kotlinTypeFacade: KotlinTypeFacade): Interpreter.InterpretationResult<PluginDataFrameSchema>
+}
+
+abstract class AbstractSchemaModificationInterpreter :
+    AbstractInterpreter<PluginDataFrameSchema>(),
+    SchemaModificationInterpreter

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/SimpleCol.kt

@@ -0,0 +1,137 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl
+
+import org.jetbrains.kotlin.fir.analysis.checkers.fullyExpandedClassId
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeNullability
+import org.jetbrains.kotlin.fir.types.isNullable
+import org.jetbrains.kotlin.fir.types.renderReadable
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.wrap
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.pluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+
+data class PluginDataFrameSchema(
+    private val columns: List<SimpleCol>
+) {
+    companion object {
+        val EMPTY = PluginDataFrameSchema(emptyList())
+    }
+
+    fun columns(): List<SimpleCol> {
+        return columns
+    }
+
+    override fun toString(): String {
+        return columns.asString()
+    }
+}
+
+fun PluginDataFrameSchema.add(name: String, type: ConeKotlinType, context: KotlinTypeFacade): PluginDataFrameSchema {
+    return PluginDataFrameSchema(columns() + context.simpleColumnOf(name, type))
+}
+
+private fun List<SimpleCol>.asString(indent: String = ""): String {
+    if (isEmpty()) return "$indent<empty compile time schema>"
+    return joinToString("\n") {
+        val col = when (it) {
+            is SimpleFrameColumn -> {
+                "${it.name}: *\n" + it.columns().asString("$indent ")
+            }
+
+            is SimpleColumnGroup -> {
+                "${it.name}:\n" + it.columns().asString("$indent ")
+            }
+
+            is SimpleDataColumn -> {
+                "${it.name}: ${it.type.type.renderReadable()}"
+            }
+        }
+        "$indent$col"
+    }
+}
+
+sealed interface SimpleCol {
+    val name: String
+
+    fun name(): String {
+        return name
+    }
+
+    fun rename(s: String): SimpleCol
+}
+
+data class SimpleDataColumn(
+    override val name: String,
+    val type: TypeApproximation
+) : SimpleCol {
+
+    override fun name(): String {
+        return name
+    }
+
+    override fun rename(s: String): SimpleDataColumn {
+        return SimpleDataColumn(s, type)
+    }
+
+    fun changeType(type: TypeApproximation): SimpleDataColumn {
+        return SimpleDataColumn(name, type)
+    }
+
+}
+
+data class SimpleFrameColumn(
+    override val name: String,
+    private val columns: List<SimpleCol>
+) : SimpleCol {
+    fun columns(): List<SimpleCol> {
+        return columns
+    }
+
+    override fun rename(s: String): SimpleFrameColumn {
+        return SimpleFrameColumn(s, columns)
+    }
+}
+
+data class SimpleColumnGroup(
+    override val name: String,
+    private val columns: List<SimpleCol>
+) : SimpleCol {
+
+    fun columns(): List<SimpleCol> {
+        return columns
+    }
+
+    override fun rename(s: String): SimpleColumnGroup {
+        return SimpleColumnGroup(s, columns)
+    }
+}
+
+fun KotlinTypeFacade.simpleColumnOf(name: String, type: ConeKotlinType): SimpleCol {
+    return if (type.fullyExpandedClassId(session) == Names.DATA_ROW_CLASS_ID) {
+        val schema = pluginDataFrameSchema(type.typeArguments[0])
+        val group = SimpleColumnGroup(name, schema.columns())
+        val column = if (type.isNullable) {
+            makeNullable(group)
+        } else {
+            group
+        }
+        column
+    } else if (type.fullyExpandedClassId(session) == Names.DF_CLASS_ID && type.nullability == ConeNullability.NOT_NULL) {
+        val schema = pluginDataFrameSchema(type.typeArguments[0])
+        SimpleFrameColumn(name, schema.columns())
+    } else {
+        SimpleDataColumn(name, type.wrap())
+    }
+}
+
+internal fun KotlinTypeFacade.makeNullable(column: SimpleCol): SimpleCol {
+    return when (column) {
+        is SimpleColumnGroup -> {
+            SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) })
+        }
+
+        is SimpleFrameColumn -> column
+        is SimpleDataColumn -> SimpleDataColumn(column.name, column.type.changeNullability { true })
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/Nulls.kt

@@ -0,0 +1,64 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlin.fir.types.ConeNullability
+import org.jetbrains.kotlin.fir.types.ConeSimpleKotlinType
+import org.jetbrains.kotlin.fir.types.typeContext
+import org.jetbrains.kotlin.fir.types.withNullability
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.Marker
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+
+class DropNulls0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return PluginDataFrameSchema(fillNullsImpl(receiver.columns(), columns.resolve(receiver).mapTo(mutableSetOf()) { it.path.path }, emptyList()))
+    }
+}
+
+class DropNa0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.whereAllNA: Boolean by arg(defaultValue = Present(false))
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        if (whereAllNA) return receiver
+        return PluginDataFrameSchema(fillNullsImpl(receiver.columns(), columns.resolve(receiver).mapTo(mutableSetOf()) { it.path.path }, emptyList()))
+    }
+}
+
+fun KotlinTypeFacade.fillNullsImpl(
+    columns: List<SimpleCol>,
+    paths: Set<List<String>>,
+    p: List<String>
+): List<SimpleCol> {
+    return columns.map {
+        // else report?
+        if (p + it.name() in paths && it is SimpleDataColumn) {
+            val coneType = it.type.type as? ConeSimpleKotlinType
+            if (coneType != null) {
+                val type = coneType.withNullability(ConeNullability.NOT_NULL, session.typeContext)
+                it.changeType(Marker.invoke(type))
+            } else {
+                // report?
+                it
+            }
+        } else {
+            if (it is SimpleColumnGroup) {
+                val updatedColumns = fillNullsImpl(it.columns(), paths, p + it.name())
+                SimpleColumnGroup(it.name(), updatedColumns)
+            } else {
+                it
+            }
+        }
+    }
+}
+

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/ReducedGroupBy.kt

@@ -0,0 +1,35 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.groupBy
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
+import org.jetbrains.kotlinx.dataframe.plugin.impl.makeNullable
+
+class GroupByReducePredicate : AbstractInterpreter<GroupBy>() {
+    val Arguments.receiver by groupBy()
+    val Arguments.predicate by ignore()
+    override fun Arguments.interpret(): GroupBy {
+        return receiver
+    }
+}
+
+class GroupByReduceExpression : AbstractInterpreter<GroupBy>() {
+    val Arguments.receiver by groupBy()
+    val Arguments.rowExpression by ignore()
+    override fun Arguments.interpret(): GroupBy {
+        return receiver
+    }
+}
+
+class GroupByReduceInto : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by groupBy()
+    val Arguments.columnName: String by arg()
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val group = makeNullable(SimpleColumnGroup(columnName, receiver.groups.columns()))
+        return PluginDataFrameSchema(receiver.keys.columns() + group)
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/add.kt

@@ -0,0 +1,84 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.Marker
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dsl
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
+import org.jetbrains.kotlinx.dataframe.plugin.impl.type
+
+typealias TypeApproximation = Marker
+
+class Add : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.name: String by arg()
+    val Arguments.infer by ignore()
+    val Arguments.type: TypeApproximation by type(name("expression"))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return PluginDataFrameSchema(receiver.columns() + simpleColumnOf(name, type.type))
+    }
+}
+
+class From : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: AddDslApproximation by arg()
+    val Arguments.receiver: String by arg()
+    val Arguments.type: TypeApproximation by type(name("expression"))
+
+    override fun Arguments.interpret() {
+        dsl.columns += simpleColumnOf(receiver, type.type)
+    }
+}
+
+class Into : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: AddDslApproximation by arg()
+    val Arguments.receiver: TypeApproximation by type()
+    val Arguments.name: String by arg()
+
+    override fun Arguments.interpret() {
+        dsl.columns += simpleColumnOf(name, receiver.type)
+    }
+}
+
+class AddDslApproximation(val columns: MutableList<SimpleCol>)
+
+class AddWithDsl : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.body by dsl()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val addDsl = AddDslApproximation(receiver.columns().toMutableList())
+        body(addDsl, emptyMap())
+        return PluginDataFrameSchema(addDsl.columns)
+    }
+}
+
+class AddDslStringInvoke : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: AddDslApproximation by arg()
+    val Arguments.receiver: String by arg()
+    val Arguments.body by dsl()
+
+    override fun Arguments.interpret() {
+        val addDsl = AddDslApproximation(mutableListOf())
+        body(addDsl, emptyMap())
+        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))
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/addId.kt

@@ -0,0 +1,21 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
+
+class AddId : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columnName: String by arg(defaultValue = Present("id"))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = buildList {
+            add(simpleColumnOf(columnName, session.builtinTypes.intType.type))
+            addAll(receiver.columns())
+        }
+        return PluginDataFrameSchema(columns)
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/columnsSelectionDsl.kt

@@ -0,0 +1,79 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.ColumnsSelector
+import org.jetbrains.kotlinx.dataframe.api.cast
+import org.jetbrains.kotlinx.dataframe.api.getColumnsWithPaths
+import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
+import org.jetbrains.kotlinx.dataframe.columns.ColumnResolutionContext
+import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
+import org.jetbrains.kotlinx.dataframe.columns.SingleColumn
+import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ConeTypesAdapter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asSimpleColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+
+internal class And10 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.other: ColumnsResolver by arg()
+    val Arguments.receiver: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return object : ColumnsResolver {
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                return receiver.resolve(df) + other.resolve(df)
+            }
+        }
+    }
+}
+
+class SingleColumnApproximation(val col: ColumnWithPathApproximation) : ColumnsResolver, SingleColumn<Any?>, ColumnReference<Any?> {
+
+    override fun name(): String {
+        return col.column.name
+    }
+
+    override fun rename(newName: String): ColumnReference<Any?> {
+        return SingleColumnApproximation(ColumnWithPathApproximation(col.path, col.column.rename(newName)))
+    }
+
+    override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+        return listOf(col)
+    }
+
+    override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<Any?>> {
+        return listOf(resolveSingle(context))
+    }
+
+    override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<Any?> {
+        return ColumnWithPath(col.column.asDataColumn(), col.path)
+    }
+}
+
+interface ColumnsResolver : org.jetbrains.kotlinx.dataframe.columns.ColumnSet<Any?> {
+    fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation>
+
+    override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<Any?>> {
+        val schema = context.df.cast<ConeTypesAdapter>().toPluginDataFrameSchema()
+        return resolve(schema).map { ColumnWithPath(it.column.asDataColumn(), it.path) }
+    }
+}
+
+abstract class ColumnsResolverAdapter : org.jetbrains.kotlinx.dataframe.columns.ColumnSet<Any?>, ColumnsResolver {
+    override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+        return resolve(ColumnResolutionContext(df.asDataFrame(), UnresolvedColumnsPolicy.Skip))
+            .map { ColumnWithPathApproximation(it.path, it.data.asSimpleColumn()) }
+    }
+}
+
+fun columnsResolver(f: ColumnsSelector<*, *>): ColumnsResolver {
+    return object : ColumnsResolverAdapter() {
+        override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<Any?>> {
+            return context.df.getColumnsWithPaths(f)
+        }
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/convert.kt

@@ -0,0 +1,200 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
+import org.jetbrains.kotlin.fir.declarations.getBooleanArgument
+import org.jetbrains.kotlin.fir.declarations.getKClassArgument
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.references.toResolvedFunctionSymbol
+import org.jetbrains.kotlin.fir.types.ConeNullability
+import org.jetbrains.kotlin.fir.types.typeContext
+import org.jetbrains.kotlin.fir.types.withNullability
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlinx.dataframe.api.asColumn
+import org.jetbrains.kotlinx.dataframe.api.convert
+import org.jetbrains.kotlinx.dataframe.api.toPath
+import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.wrap
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Absent
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.type
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+
+internal class Convert0 : AbstractInterpreter<ConvertApproximation>() {
+    val Arguments.columns: ColumnsResolver by arg()
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    override val Arguments.startingSchema get() = receiver
+
+    override fun Arguments.interpret(): ConvertApproximation {
+        return ConvertApproximation(receiver, columns.resolve(receiver).map { it.path.path })
+    }
+}
+
+class Convert2 : AbstractInterpreter<ConvertApproximation>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: List<String> by arg(defaultValue = Absent)
+
+    override fun Arguments.interpret(): ConvertApproximation {
+        return ConvertApproximation(receiver, columns.map { listOf(it) })
+    }
+}
+
+class ConvertApproximation(val schema: PluginDataFrameSchema, val columns: List<List<String>>)
+
+internal class Convert6 : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.firstCol: String by arg()
+    val Arguments.cols: List<String> by arg(defaultValue = Present(emptyList()))
+    val Arguments.infer by ignore()
+    val Arguments.expression: TypeApproximation by type()
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    override val Arguments.startingSchema get() = receiver
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = (listOf(firstCol) + cols).map { listOf(it) }
+        return convertImpl(receiver, columns, expression)
+    }
+}
+
+class With0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: ConvertApproximation by arg()
+    val Arguments.infer by ignore()
+    val Arguments.type: TypeApproximation by type(name("rowConverter"))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return convertImpl(receiver.schema, receiver.columns, type)
+    }
+}
+
+class PerRowCol : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: ConvertApproximation by arg()
+    val Arguments.infer by ignore()
+    val Arguments.type: TypeApproximation by type(name("expression"))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return convertImpl(receiver.schema, receiver.columns, type)
+    }
+}
+
+internal fun KotlinTypeFacade.convertImpl(
+    pluginDataFrameSchema: PluginDataFrameSchema,
+    columns: List<List<String>>,
+    type: TypeApproximation
+): PluginDataFrameSchema {
+    return pluginDataFrameSchema.map(columns.toSet()) { path, column ->
+        val unwrappedType = type.type
+        simpleColumnOf(column.name, unwrappedType)
+    }
+}
+
+internal fun PluginDataFrameSchema.map(selected: ColumnsSet, transform: ColumnMapper): PluginDataFrameSchema {
+    return PluginDataFrameSchema(
+        f(columns(), transform, selected, emptyList())
+    )
+}
+
+internal typealias ColumnsSet = Set<List<String>>
+
+internal typealias ColumnMapper = (List<String>, SimpleCol) -> SimpleCol
+
+internal fun f(columns: List<SimpleCol>, transform: ColumnMapper, selected: ColumnsSet, path: List<String>): List<SimpleCol> {
+    return columns.map {
+        val fullPath = path + listOf(it.name)
+        when (it) {
+            is SimpleColumnGroup -> if (fullPath in selected) {
+                transform(fullPath, it)
+            } else {
+                it.map(transform, selected, fullPath)
+            }
+            is SimpleFrameColumn -> if (fullPath in selected) {
+                transform(fullPath, it)
+            } else {
+                it.map(transform, selected, fullPath)
+            }
+            is SimpleDataColumn -> if (fullPath in selected) {
+                transform(path, it)
+            } else {
+                it
+            }
+        }
+    }
+}
+
+internal fun SimpleColumnGroup.map(transform: ColumnMapper, selected: ColumnsSet, path: List<String>): SimpleColumnGroup {
+    return SimpleColumnGroup(
+        name,
+        f(columns(), transform, selected, path)
+    )
+}
+
+internal fun SimpleFrameColumn.map(transform: ColumnMapper, selected: ColumnsSet, path: List<String>): SimpleFrameColumn {
+    return SimpleFrameColumn(
+        name,
+        f(columns(), transform, selected, path)
+    )
+}
+
+internal class To0 : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.receiver: ConvertApproximation by arg()
+    val Arguments.typeArg0: TypeApproximation by arg()
+    override val Arguments.startingSchema get() = receiver.schema
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return convertImpl(receiver.schema, receiver.columns, typeArg0)
+    }
+}
+
+internal class ConvertAsColumn : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: ConvertApproximation by arg()
+    val Arguments.typeArg2: TypeApproximation by arg()
+    val Arguments.type: TypeApproximation by type(name("columnConverter"))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.schema.asDataFrame()
+            .convert { receiver.columns.map { it.toPath() }.toColumnSet() }
+            .asColumn { simpleColumnOf("", typeArg2.type).asDataColumn() }
+            .toPluginDataFrameSchema()
+    }
+}
+
+internal abstract class AbstractToSpecificType : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.functionCall: FirFunctionCall by arg(lens = Interpreter.Id)
+    val Arguments.receiver: ConvertApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val converterAnnotation = functionCall.calleeReference.toResolvedFunctionSymbol()?.getAnnotationByClassId(Names.CONVERTER_ANNOTATION, session)
+        val to = converterAnnotation?.getKClassArgument(Name.identifier("klass"), session)
+        val nullable = converterAnnotation?.getBooleanArgument(Name.identifier("nullable"), session)
+        return if (to != null && nullable != null) {
+            val targetType = to.withNullability(ConeNullability.create(nullable), session.typeContext)
+            convertImpl(receiver.schema, receiver.columns, targetType.wrap())
+        } else {
+            PluginDataFrameSchema.EMPTY
+        }
+    }
+}
+
+internal class ToSpecificType : AbstractToSpecificType()
+
+internal class ToSpecificTypeZone : AbstractToSpecificType() {
+    val Arguments.zone by ignore()
+}
+
+internal class ToSpecificTypePattern : AbstractToSpecificType() {
+    val Arguments.pattern by ignore()
+    val Arguments.locale by ignore()
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/count.kt

@@ -0,0 +1,19 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.add
+import org.jetbrains.kotlinx.dataframe.plugin.impl.groupBy
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
+
+class GroupByCount0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by groupBy()
+    val Arguments.resultName: String by arg(defaultValue = Present("count"))
+    val Arguments.predicate by ignore()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.keys.add(resultName, session.builtinTypes.intType.type, context = this)
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/dataFrameOf.kt

@@ -0,0 +1,85 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
+import org.jetbrains.kotlin.fir.expressions.FirVarargArgumentsExpression
+import org.jetbrains.kotlin.fir.plugin.createConeType
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.commonSuperTypeOrNull
+import org.jetbrains.kotlin.fir.types.resolvedType
+import org.jetbrains.kotlin.fir.types.type
+import org.jetbrains.kotlin.fir.types.typeContext
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
+import org.jetbrains.kotlinx.dataframe.impl.api.withValuesImpl
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+
+class DataFrameOf0 : AbstractInterpreter<DataFrameBuilderApproximation>() {
+    val Arguments.header: List<String> by arg()
+
+    override fun Arguments.interpret(): DataFrameBuilderApproximation {
+        return DataFrameBuilderApproximation(header)
+    }
+}
+
+class DataFrameBuilderApproximation(val header: List<String>)
+
+class DataFrameBuilderInvoke0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: DataFrameBuilderApproximation by arg()
+    val Arguments.values: FirVarargArgumentsExpression by arg(lens = Interpreter.Id)
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = (receiver.header to values.arguments).withValuesImpl().map { (name, values) ->
+            val type = session.typeContext.commonSuperTypeOrNull(values.map { it.resolvedType }) ?: error("$name $values")
+            simpleColumnOf(name, type)
+        }
+        return PluginDataFrameSchema(columns)
+    }
+}
+
+class DataFrameOf3 : AbstractSchemaModificationInterpreter() {
+    val Arguments.columns: List<Interpreter.Success<Pair<*, *>>> by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val res = columns.map {
+            val it = it.value
+            val name = (it.first as? FirLiteralExpression)?.value as? String
+            val type = (it.second as? FirExpression)?.resolvedType?.typeArguments?.getOrNull(0)?.type
+            if (name == null || type == null) return PluginDataFrameSchema(emptyList())
+            simpleColumnOf(name, type)
+        }
+        return PluginDataFrameSchema(res)
+    }
+}
+
+abstract class SchemaConstructor : AbstractSchemaModificationInterpreter() {
+    val Arguments.columns: List<Interpreter.Success<Pair<*, *>>> by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val res = columns.map {
+            val it = it.value
+            val name = (it.first as? FirLiteralExpression)?.value as? String
+            val resolvedType = (it.second as? FirExpression)?.resolvedType
+            val type: ConeKotlinType? = when (resolvedType?.classId) {
+                Names.COLUM_GROUP_CLASS_ID -> Names.DATA_ROW_CLASS_ID.createConeType(session, arrayOf(resolvedType.typeArguments[0]))
+                Names.FRAME_COLUMN_CLASS_ID -> Names.DF_CLASS_ID.createConeType(session, arrayOf(resolvedType.typeArguments[0]))
+                Names.DATA_COLUMN_CLASS_ID -> resolvedType.typeArguments[0] as? ConeKotlinType
+                Names.BASE_COLUMN_CLASS_ID -> resolvedType.typeArguments[0] as? ConeKotlinType
+                else -> null
+            }
+            if (name == null || type == null) return PluginDataFrameSchema(emptyList())
+            simpleColumnOf(name, type)
+        }
+        return PluginDataFrameSchema(res)
+    }
+}
+
+class DataFrameOfPairs : SchemaConstructor()
+
+class ColumnOfPairs : SchemaConstructor()
+

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/explode.kt

@@ -0,0 +1,85 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
+
+internal class Explode0 : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.dropEmpty: Boolean by arg(defaultValue = Present(true))
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: ColumnsResolver? by arg(defaultValue = Present(null))
+    override val Arguments.startingSchema get() = receiver
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = columns ?: object : ColumnsResolver {
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                return df.flatten(includeFrames = false).filter {
+                    val column = it.column
+                    column is SimpleFrameColumn || column is SimpleDataColumn && column.type.isList()
+                }
+            }
+        }
+        return receiver.explodeImpl(dropEmpty, columns.resolve(receiver))
+    }
+}
+
+val KotlinTypeFacade.explodeImpl: PluginDataFrameSchema.(dropEmpty: Boolean, selector: List<ColumnWithPathApproximation>) -> PluginDataFrameSchema
+    get()  = { dropEmpty, selector ->
+    val columns = selector
+
+    val selected = columns.associateBy { it.path }
+
+    fun makeNullable(column: SimpleCol): SimpleCol {
+        return when (column) {
+            is SimpleColumnGroup -> SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) })
+            is SimpleFrameColumn -> column
+            is SimpleDataColumn -> {
+                column.changeType(type = column.type.changeNullability { nullable -> selector.size > 1 || !dropEmpty || nullable })
+            }
+        }
+    }
+
+    fun explode(column: SimpleCol, path: List<String>): SimpleCol {
+        val fullPath = path + listOf(column.name)
+        return when (column) {
+            is SimpleColumnGroup -> {
+                SimpleColumnGroup(column.name, column.columns().map { explode(it, fullPath) })
+            }
+            is SimpleFrameColumn -> {
+                val s = selected[fullPath]
+                if (s != null) {
+                    SimpleColumnGroup(s.column.name, column.columns().map { makeNullable(it) })
+                } else {
+                    column
+                }
+            }
+            is SimpleDataColumn -> {
+                val s = selected[fullPath]
+                if (s != null) {
+                    val newType = when {
+                        column.type.isList() -> column.type.typeArgument()
+                        else -> column.type
+                    }
+                    makeNullable(simpleColumnOf(s.column.name, newType.type))
+                } else {
+                    column
+                }
+            }
+        }
+    }
+
+    PluginDataFrameSchema(
+        columns().map { column ->
+            explode(column, emptyList())
+        }
+    )
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/fillNulls.kt

@@ -0,0 +1,31 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.type
+
+class FillNulls0 : AbstractInterpreter<FillNullsApproximation>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): FillNullsApproximation {
+        return FillNullsApproximation(receiver, columns)
+    }
+}
+
+class FillNullsApproximation(val schema: PluginDataFrameSchema, val columns: ColumnsResolver) : UpdateApproximation
+
+class UpdateWith0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: UpdateApproximation by arg()
+    val Arguments.expression: TypeApproximation by type()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return when (val receiver = receiver) {
+            is FillNullsApproximation -> convertImpl(receiver.schema, receiver.columns.resolve(receiver.schema).map { it.path.path }, expression)
+            is UpdateApproximationImpl -> convertImpl(receiver.schema, receiver.columns.resolve(receiver.schema).map { it.path.path }, expression)
+        }
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/flatten.kt

@@ -0,0 +1,38 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.api.flatten
+import org.jetbrains.kotlinx.dataframe.api.pathOf
+import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+
+class FlattenDefault : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by dataFrame()
+    val Arguments.keepParentNameForColumns: Boolean by arg(defaultValue = Present(false))
+    val Arguments.separator: String by arg(defaultValue = Present("_"))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.asDataFrame().flatten(keepParentNameForColumns, separator).toPluginDataFrameSchema()
+    }
+}
+
+class Flatten0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by dataFrame()
+    val Arguments.keepParentNameForColumns: Boolean by arg(defaultValue = Present(false))
+    val Arguments.separator: String by arg(defaultValue = Present("_"))
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = columns.resolve(receiver).map { it.path }
+        return receiver
+            .asDataFrame()
+            .flatten(keepParentNameForColumns, separator) { columns.toColumnSet() }
+            .toPluginDataFrameSchema()
+    }
+}
+

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/flattenPluginSchema.kt

@@ -0,0 +1,35 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
+
+fun PluginDataFrameSchema.flatten(includeFrames: Boolean): List<ColumnWithPathApproximation> {
+    if (columns().isEmpty()) return emptyList()
+    val columns = mutableListOf<ColumnWithPathApproximation>()
+    flattenImpl(columns(), emptyList(), columns, includeFrames)
+    return columns
+}
+
+fun flattenImpl(columns: List<SimpleCol>, path: List<String>, flatList: MutableList<ColumnWithPathApproximation>, includeFrames: Boolean) {
+    columns.forEach { column ->
+        val fullPath = path + listOf(column.name)
+        when (column) {
+            is SimpleColumnGroup -> {
+                flatList.add(ColumnWithPathApproximation(ColumnPathApproximation(fullPath), column))
+                flattenImpl(column.columns(), fullPath, flatList, includeFrames)
+            }
+            is SimpleFrameColumn -> {
+                flatList.add(ColumnWithPathApproximation(ColumnPathApproximation(fullPath), column))
+                flattenImpl(column.columns(), fullPath, flatList, includeFrames)
+            }
+            is SimpleDataColumn -> {
+                flatList.add(ColumnWithPathApproximation(ColumnPathApproximation(fullPath), column))
+            }
+        }
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/group.kt

@@ -0,0 +1,31 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.api.group
+import org.jetbrains.kotlinx.dataframe.api.into
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+
+class Group0 : AbstractInterpreter<GroupClauseApproximation>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): GroupClauseApproximation {
+        return GroupClauseApproximation(receiver, columns)
+    }
+}
+
+class GroupClauseApproximation(val df: PluginDataFrameSchema, val columns: ColumnsResolver)
+
+class Into0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: GroupClauseApproximation by arg()
+    val Arguments.column: String by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.df.asDataFrame().group { receiver.columns }.into(column).toPluginDataFrameSchema()
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/groupBy.kt

@@ -0,0 +1,421 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeNullability
+import org.jetbrains.kotlin.fir.types.constructClassLikeType
+import org.jetbrains.kotlin.fir.types.isNullable
+import org.jetbrains.kotlin.fir.types.isSubtypeOf
+import org.jetbrains.kotlin.fir.types.resolvedType
+import org.jetbrains.kotlin.fir.types.typeContext
+import org.jetbrains.kotlin.fir.types.withNullability
+import org.jetbrains.kotlin.name.StandardClassIds
+import org.jetbrains.kotlinx.dataframe.api.remove
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator
+import org.jetbrains.kotlinx.dataframe.plugin.InterpretationErrorReporter
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.add
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.groupBy
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
+import org.jetbrains.kotlinx.dataframe.plugin.impl.makeNullable
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.type
+import org.jetbrains.kotlinx.dataframe.plugin.interpret
+import org.jetbrains.kotlinx.dataframe.plugin.loadInterpreter
+import kotlin.collections.plus
+
+class GroupBy(val keys: PluginDataFrameSchema, val groups: PluginDataFrameSchema) {
+    companion object {
+        val EMPTY = GroupBy(PluginDataFrameSchema.EMPTY, PluginDataFrameSchema.EMPTY)
+    }
+}
+
+class DataFrameGroupBy : AbstractInterpreter<GroupBy>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.moveToTop: Boolean by arg(defaultValue = Present(true))
+    val Arguments.cols: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): GroupBy {
+        return GroupBy(keys = createPluginDataFrameSchema(cols.resolve(receiver), moveToTop), groups = receiver)
+    }
+}
+
+class AsGroupBy : AbstractInterpreter<GroupBy>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.selector: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): GroupBy {
+        val column = selector.resolve(receiver).singleOrNull()?.column
+        return if (column is SimpleFrameColumn) {
+            GroupBy(receiver.asDataFrame().remove { selector }.toPluginDataFrameSchema(), PluginDataFrameSchema(column.columns()))
+        } else {
+            GroupBy.EMPTY
+        }
+    }
+}
+
+class AsGroupByDefault : AbstractInterpreter<GroupBy>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+
+    override fun Arguments.interpret(): GroupBy {
+        val groups = receiver.columns().singleOrNull { it is SimpleFrameColumn } as? SimpleFrameColumn
+        return if (groups != null) {
+            GroupBy(receiver.asDataFrame().remove(groups.name).toPluginDataFrameSchema(), PluginDataFrameSchema(groups.columns()))
+        } else {
+            GroupBy.EMPTY
+        }
+    }
+}
+
+class NamedValue(val name: String, val type: ConeKotlinType)
+
+class GroupByDsl {
+    val columns = mutableListOf<NamedValue>()
+}
+
+class AggregateDslInto : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: GroupByDsl by arg()
+    val Arguments.receiver: FirExpression by arg(lens = Interpreter.Id)
+    val Arguments.name: String by arg()
+
+    override fun Arguments.interpret() {
+        dsl.columns.add(NamedValue(name, receiver.resolvedType))
+    }
+}
+
+class Aggregate : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: GroupBy by groupBy()
+    val Arguments.body: FirAnonymousFunctionExpression by arg(lens = Interpreter.Id)
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return aggregate(
+            receiver,
+            InterpretationErrorReporter.DEFAULT,
+            body
+        )
+    }
+}
+
+class AggregateRow : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.body: FirAnonymousFunctionExpression by arg(lens = Interpreter.Id)
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return aggregate(
+            GroupBy(PluginDataFrameSchema.EMPTY, receiver),
+            InterpretationErrorReporter.DEFAULT,
+            body
+        )
+    }
+}
+
+fun KotlinTypeFacade.aggregate(
+    groupBy: GroupBy,
+    reporter: InterpretationErrorReporter,
+    firAnonymousFunctionExpression: FirAnonymousFunctionExpression
+): PluginDataFrameSchema {
+    val body = firAnonymousFunctionExpression.anonymousFunction.body
+    val lastExpression = (body?.statements?.lastOrNull() as? FirReturnExpression)?.result
+    val type = lastExpression?.resolvedType
+    return if (type != session.builtinTypes.unitType) {
+        val dsl = GroupByDsl()
+        val calls = buildList {
+            body?.statements?.filterIsInstance<FirFunctionCall>()?.let { addAll(it) }
+            if (lastExpression is FirFunctionCall) add(lastExpression)
+        }
+        calls.forEach { call ->
+            val schemaProcessor = call.loadInterpreter() ?: return@forEach
+            interpret(
+                call,
+                schemaProcessor,
+                mapOf("dsl" to Interpreter.Success(dsl)),
+                reporter
+            )
+        }
+
+        val cols = groupBy.keys.columns() + dsl.columns.map {
+            simpleColumnOf(it.name, it.type)
+        }
+        PluginDataFrameSchema(cols)
+    } else {
+        PluginDataFrameSchema.EMPTY
+    }
+}
+
+fun KotlinTypeFacade.createPluginDataFrameSchema(
+    keys: List<ColumnWithPathApproximation>,
+    moveToTop: Boolean
+): PluginDataFrameSchema {
+    fun addToHierarchy(
+        path: List<String>,
+        column: SimpleCol,
+        columns: List<SimpleCol>
+    ): List<SimpleCol> {
+        if (path.isEmpty()) return columns
+
+        val groupName = path[0]
+        val remainingPath = path.drop(1)
+
+        val updatedColumns = columns.map {
+            if (it is SimpleColumnGroup && it.name == groupName) {
+                SimpleColumnGroup(it.name, columns = addToHierarchy(remainingPath, column, it.columns()))
+            } else {
+                it
+            }
+        }
+
+        return if (updatedColumns.any { it is SimpleColumnGroup && it.name == groupName }) {
+            updatedColumns
+        } else {
+            val newGroup = if (remainingPath.isEmpty()) {
+                column
+            } else {
+                SimpleColumnGroup(groupName, addToHierarchy(remainingPath, column, emptyList()))
+            }
+            updatedColumns + newGroup
+        }
+    }
+
+    var rootColumns: List<SimpleCol> = emptyList()
+
+    if (moveToTop) {
+        rootColumns = keys.map { it.column }
+    } else {
+        for (key in keys) {
+            val path = key.path.path
+            val column = key.column
+            rootColumns = addToHierarchy(path, column, rootColumns)
+        }
+    }
+
+
+    return PluginDataFrameSchema(rootColumns)
+}
+
+class GroupByInto : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: GroupBy by groupBy()
+    val Arguments.column: String by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val grouped = listOf(SimpleFrameColumn(column, receiver.groups.columns()))
+        return PluginDataFrameSchema(
+            receiver.keys.columns() + grouped
+        )
+    }
+}
+
+class GroupByToDataFrame : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: GroupBy by groupBy()
+    val Arguments.groupedColumnName: String? by arg(defaultValue = Present(null))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val grouped = listOf(SimpleFrameColumn(groupedColumnName ?: "group", receiver.groups.columns()))
+        return PluginDataFrameSchema(
+            receiver.keys.columns() + grouped
+        )
+    }
+}
+
+class GroupByAdd : AbstractInterpreter<GroupBy>() {
+    val Arguments.receiver: GroupBy by groupBy()
+    val Arguments.name: String by arg()
+    val Arguments.infer by ignore()
+    val Arguments.type: TypeApproximation by type(name("expression"))
+
+    override fun Arguments.interpret(): GroupBy {
+        return GroupBy(receiver.keys, receiver.groups.add(name, type.type, context = this))
+    }
+}
+
+/** Produces a type of aggregated column based on the expression type. */
+abstract class GroupByAggregatorOf(val defaultName: String, val aggregator: Aggregator<*, *>) : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by groupBy()
+    val Arguments.name: String? by arg(defaultValue = Present(null))
+    val Arguments.expression by type()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val aggregated = makeNullable(simpleColumnOf(name ?: defaultName, expression.type))
+
+        val newColumns = generateStatisticResultColumns(aggregator, listOf(aggregated as SimpleDataColumn))
+        return PluginDataFrameSchema(receiver.keys.columns() + newColumns)
+    }
+}
+
+/** Implementation for `maxOf`. */
+class GroupByMaxOf : GroupByAggregatorOf(defaultName = "max", max)
+
+/** Implementation for `minOf`. */
+class GroupByMinOf : GroupByAggregatorOf(defaultName = "min", min)
+
+/** Implementation for `medianOf`. */
+class GroupByMedianOf : GroupByAggregatorOf(defaultName = "median", median)
+
+/** Implementation for `meanOf`. */
+class GroupByMeanOf : GroupByAggregatorOf(defaultName = "mean", mean)
+
+/** Implementation for `stdOf`. */
+class GroupByStdOf : GroupByAggregatorOf(defaultName = "std", std)
+
+/**
+ * Provides a base implementation for a custom schema modification interpreter
+ * that groups data by specified criteria and produces aggregated results.
+ *
+ * The class uses a `defaultName` to define a fallback name for the result column
+ * if no specific name is provided. It leverages `Arguments` properties to define
+ * and resolve the group-by receiver, result name, and expression type.
+ *
+ * Key Components:
+ * - [receiver] Represents the input data that will be grouped.
+ * - [resultName] Optional name for the resulting aggregated column. Defaults to `defaultName`.
+ * - [expression] Defines the type of the expression for aggregation.
+ */
+abstract class GroupByAggregatorSumOf(val defaultName: String) : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by groupBy()
+    val Arguments.resultName: String? by arg(defaultValue = Present(null))
+    val Arguments.expression by type()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val aggregated = makeNullable(simpleColumnOf(resultName ?: defaultName, expression.type))
+        val newColumns = generateStatisticResultColumns(sum, listOf(aggregated as SimpleDataColumn))
+        return PluginDataFrameSchema(receiver.keys.columns() + newColumns)
+    }
+}
+
+/** Implementation for `sumOf`. */
+class GroupBySumOf : GroupByAggregatorSumOf(defaultName = "sum")
+
+/**
+ * Provides a base implementation for a custom schema modification interpreter
+ * that groups data by specified criteria and produces aggregated results.
+ *
+ * The class uses a `defaultName` to define a fallback name for the result column
+ * if no specific name is provided. It leverages `Arguments` properties to define
+ * and resolve the group-by receiver, result name, and expression type.
+ *
+ * Key Components:
+ * - [receiver] Represents the input data that will be grouped.
+ * - [name] Optional name for the resulting aggregated column. Defaults to `defaultName`.
+ * - [columns] ColumnsResolver to define which columns to include in the grouping operation.
+ */
+abstract class GroupByAggregator0(val defaultName: String, val aggregator: Aggregator<*, *>) : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by groupBy()
+    val Arguments.name: String? by arg(defaultValue = Present(null))
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        if (name == null) {
+            val resolvedColumns = columns.resolve(receiver.keys).map { (it.column as SimpleDataColumn) }.toList()
+            val newColumns = generateStatisticResultColumns(aggregator, resolvedColumns)
+
+            return PluginDataFrameSchema(receiver.keys.columns() + newColumns)
+        } else {
+            val resolvedColumns = columns.resolve(receiver.keys).map { it.column }.toList()
+            // TODO: handle multiple columns https://github.com/Kotlin/dataframe/issues/1090
+            val aggregated =
+                makeNullable(simpleColumnOf(name ?: defaultName, (resolvedColumns[0] as SimpleDataColumn).type.type))
+            // case with the multiple columns will be removed for GroupBy
+            val newColumns = generateStatisticResultColumns(aggregator, listOf(aggregated as SimpleDataColumn))
+            return PluginDataFrameSchema(receiver.keys.columns() + newColumns)
+        }
+    }
+}
+
+/** Implementation for `sum`. */
+class GroupBySum0 : GroupByAggregator0(defaultName = "sum", sum)
+
+/** Implementation for `median`. */
+class GroupByMedian0 : GroupByAggregator0(defaultName = "median", median)
+
+/** Implementation for `median`. */
+class GroupByMin0 : GroupByAggregator0(defaultName = "min", min)
+
+/** Implementation for `median`. */
+class GroupByMax0 : GroupByAggregator0(defaultName = "max", max)
+
+/** Implementation for `mean`. */
+class GroupByMean0 : GroupByAggregator0(defaultName = "mean", mean)
+
+/** Implementation for `std`. */
+class GroupByStd0 : GroupByAggregator0(defaultName = "std", std)
+
+/** Adds to the schema only numerical columns. */
+abstract class GroupByAggregator1(val aggregator: Aggregator<*, *>) : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by groupBy()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val resolvedColumns = receiver.groups.columns()
+            .filterIsInstance<SimpleDataColumn>()
+            .filter { it.type.type.isSubtypeOf(session.builtinTypes.numberType.type, session) }
+
+        val newColumns = generateStatisticResultColumns(aggregator, resolvedColumns)
+
+        return PluginDataFrameSchema(receiver.keys.columns() + newColumns)
+    }
+}
+
+/** Implementation for `sum`. */
+class GroupBySum1 : GroupByAggregator1(sum)
+
+/** Implementation for `mean`. */
+class GroupByMean1 : GroupByAggregator1(mean)
+
+/** Implementation for `std`. */
+class GroupByStd1 : GroupByAggregator1(std)
+
+/** Keeps in schema only columns with intraComparable values. */
+abstract class GroupByAggregatorComparable(val aggregator: Aggregator<*, *>)  : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by groupBy()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val comparableColumns = receiver.groups.columns()
+            .filterIsInstance<SimpleDataColumn>()
+            .filter { isIntraComparable(it, session) }
+
+        val newColumns = generateStatisticResultColumns(aggregator, comparableColumns)
+
+        return PluginDataFrameSchema(receiver.keys.columns() + newColumns)
+    }
+}
+
+/** Implementation for `max`. */
+class GroupByMax1 : GroupByAggregatorComparable(max)
+
+/** Implementation for `min`. */
+class GroupByMin1 : GroupByAggregatorComparable(min)
+
+/** Implementation for `median`. */
+class GroupByMedian1 : GroupByAggregatorComparable(median)
+
+internal fun isIntraComparable(col: SimpleDataColumn, session: FirSession): Boolean {
+    val comparable = StandardClassIds.Comparable.constructClassLikeType(
+        typeArguments = arrayOf(col.type.type.withNullability(ConeNullability.NOT_NULL, session.typeContext)),
+        isNullable = col.type.type.isNullable,
+    )
+    return col.type.type.isSubtypeOf(comparable, session)
+}
+
+class ConcatWithKeys : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by groupBy()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val originalColumns = receiver.groups.columns()
+        return PluginDataFrameSchema(
+            originalColumns + receiver.keys.columns().filter { it.name !in originalColumns.map { it.name } })
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/insert.kt

@@ -0,0 +1,100 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.api.after
+import org.jetbrains.kotlinx.dataframe.api.at
+import org.jetbrains.kotlinx.dataframe.api.insert
+import org.jetbrains.kotlinx.dataframe.api.pathOf
+import org.jetbrains.kotlinx.dataframe.api.under
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.InsertClauseApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.type
+
+/**
+ * @see DataFrame.insert
+ */
+internal class Insert0 : AbstractInterpreter<InsertClauseApproximation>() {
+    val Arguments.receiver by dataFrame()
+    val Arguments.name: String by arg()
+    val Arguments.typeArg1 by type()
+
+    override fun Arguments.interpret(): InsertClauseApproximation {
+        return InsertClauseApproximation(receiver, simpleColumnOf(name, typeArg1.type))
+    }
+}
+
+internal class Insert1 : AbstractInterpreter<InsertClauseApproximation>() {
+    val Arguments.name: String by arg()
+    val Arguments.infer by ignore()
+    val Arguments.expression: TypeApproximation by type()
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+
+    override fun Arguments.interpret(): InsertClauseApproximation {
+        return InsertClauseApproximation(receiver, simpleColumnOf(name, expression.type))
+    }
+}
+
+internal class Under0 : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.column: SingleColumnApproximation by arg()
+    val Arguments.receiver: InsertClauseApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val path = column.resolve(receiver.df).single().path
+        return receiver.df.asDataFrame()
+            .insert(receiver.column.asDataColumn()).under(path)
+            .toPluginDataFrameSchema()
+    }
+}
+
+internal class Under1 : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.columnPath: ColumnPathApproximation by arg()
+    val Arguments.receiver: InsertClauseApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.df.asDataFrame()
+            .insert(receiver.column.asDataColumn()).under(columnPath)
+            .toPluginDataFrameSchema()
+    }
+}
+
+internal class Under4 : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.column: String by arg()
+    val Arguments.receiver: InsertClauseApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.df.asDataFrame()
+            .insert(receiver.column.asDataColumn()).under(pathOf(column))
+            .toPluginDataFrameSchema()
+    }
+}
+
+internal class InsertAfter0 : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.column: SingleColumnApproximation by arg()
+    val Arguments.receiver: InsertClauseApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.df.asDataFrame()
+                    .insert(receiver.column.asDataColumn()).after { column.col.path }
+            .toPluginDataFrameSchema()
+    }
+}
+
+internal class InsertAt : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.receiver: InsertClauseApproximation by arg()
+    val Arguments.position: Int by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.df.asDataFrame()
+            .insert(receiver.column.asDataColumn()).at(position)
+            .toPluginDataFrameSchema()
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/join.kt

@@ -0,0 +1,115 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.ColumnsContainer
+import org.jetbrains.kotlinx.dataframe.api.JoinColumnsSelector
+import org.jetbrains.kotlinx.dataframe.api.JoinDsl
+import org.jetbrains.kotlinx.dataframe.api.JoinType
+import org.jetbrains.kotlinx.dataframe.api.getColumnsWithPaths
+import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
+import org.jetbrains.kotlinx.dataframe.api.remove
+import org.jetbrains.kotlinx.dataframe.api.toDataFrameFromPairs
+import org.jetbrains.kotlinx.dataframe.columns.ColumnSet
+import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ConeTypesAdapter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asSimpleColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.enum
+import org.jetbrains.kotlinx.dataframe.plugin.impl.makeNullable
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+
+internal abstract class AbstractJoin() : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.other: PluginDataFrameSchema by dataFrame()
+    val Arguments.selector: ColumnSet<*>? by arg()
+
+    fun Arguments.join(type: JoinType): PluginDataFrameSchema {
+        val leftDf = receiver.asDataFrame()
+        val rightDf = other.asDataFrame()
+        val joinSelector: JoinColumnsSelector<*, *> = if (selector != null) {
+            { it: ColumnsContainer<*> -> selector!! }
+        } else {
+            JoinDsl.defaultJoinColumns(leftDf, rightDf)
+        }
+        val joinColumns = JoinDsl.getColumns(leftDf, rightDf, joinSelector)
+
+        val filteredRightDf = rightDf.remove { joinColumns.map { it.right }.toColumnSet() }
+        val left = leftDf.getColumnsWithPaths { colsAtAnyDepth().filter { !it.isColumnGroup() } }
+            .map { it.path to it.data }
+        val right = filteredRightDf.getColumnsWithPaths { colsAtAnyDepth().filter { !it.isColumnGroup() } }
+            .map { it.path to it.data }
+
+        val result = buildList {
+            when (type) {
+                JoinType.Inner -> {
+                    addAll(left)
+                    addAll(right)
+                }
+                JoinType.Left -> {
+                    addAll(left)
+                    addAll(right.map { it.first to makeNullable(it.second.asSimpleColumn()).asDataColumn() })
+                }
+                JoinType.Right -> {
+                    addAll(left.map { it.first to makeNullable(it.second.asSimpleColumn()).asDataColumn() })
+                    addAll(right)
+                }
+                JoinType.Full -> {
+                    addAll(left.map { it.first to makeNullable(it.second.asSimpleColumn()).asDataColumn() })
+                    addAll(right.map { it.first to makeNullable(it.second.asSimpleColumn()).asDataColumn() })
+                }
+                JoinType.Filter -> addAll(left)
+                JoinType.Exclude -> addAll(left)
+            }
+        }
+        return result.toDataFrameFromPairs<ConeTypesAdapter>().toPluginDataFrameSchema()
+    }
+}
+
+internal class Join0 : AbstractJoin() {
+    val Arguments.type: JoinType by enum(defaultValue = Present(JoinType.Inner))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(type)
+    }
+}
+
+internal class LeftJoin : AbstractJoin() {
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(JoinType.Left)
+    }
+}
+
+internal class RightJoin : AbstractJoin() {
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(JoinType.Right)
+    }
+}
+
+internal class FullJoin : AbstractJoin() {
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(JoinType.Full)
+    }
+}
+
+internal class InnerJoin : AbstractJoin() {
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(JoinType.Inner)
+    }
+}
+
+internal class FilterJoin : AbstractJoin() {
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(JoinType.Filter)
+    }
+}
+
+internal class ExcludeJoin : AbstractJoin() {
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(JoinType.Exclude)
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/joinDsl.kt

@@ -0,0 +1,31 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.api.ColumnMatch
+import org.jetbrains.kotlinx.dataframe.columns.ColumnResolutionContext
+import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
+
+internal data class ColumnMatchApproximation(
+    override val left: SingleColumnApproximation,
+    override val right: SingleColumnApproximation
+): ColumnMatch<Any?>, ColumnsResolver {
+    override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+        throw UnsupportedOperationException()
+    }
+
+    override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<Any?>> {
+        throw UnsupportedOperationException()
+    }
+}
+
+internal class Match0 : AbstractInterpreter<ColumnMatchApproximation>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.other: SingleColumnApproximation by arg()
+
+    override fun Arguments.interpret(): ColumnMatchApproximation {
+        return ColumnMatchApproximation(receiver, other)
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/joinWith.kt

@@ -0,0 +1,105 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.api.JoinType
+import org.jetbrains.kotlinx.dataframe.impl.ColumnNameGenerator
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.enum
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
+import org.jetbrains.kotlinx.dataframe.plugin.impl.makeNullable
+
+internal abstract class AbstractJoinWith() : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.right: PluginDataFrameSchema by dataFrame()
+    val Arguments.joinExpression by ignore()
+
+    fun Arguments.join(type: JoinType): PluginDataFrameSchema {
+        val left = receiver.columns()
+        val right = right.columns()
+
+        val nameGenerator = ColumnNameGenerator()
+
+        fun MutableList<SimpleCol>.addColumns(columns: List<SimpleCol>) {
+            for (column in columns) {
+                val uniqueName = nameGenerator.addUnique(column.name)
+                add(column.rename(uniqueName))
+            }
+        }
+
+        val result = buildList {
+            when (type) {
+                JoinType.Inner -> {
+                    addColumns(left)
+                    addColumns(right)
+                }
+
+                JoinType.Left -> {
+                    addColumns(left)
+                    addColumns(right.map { makeNullable(it) })
+                }
+
+                JoinType.Right -> {
+                    addColumns(left.map { makeNullable(it) })
+                    addColumns(right)
+                }
+
+                JoinType.Full -> {
+                    addColumns(left.map { makeNullable(it) })
+                    addColumns(right.map { makeNullable(it) })
+                }
+
+                JoinType.Filter -> addColumns(left)
+                JoinType.Exclude -> addColumns(left)
+            }
+        }
+        return PluginDataFrameSchema(result)
+    }
+}
+
+internal class JoinWith : AbstractJoinWith() {
+    val Arguments.type: JoinType by enum(defaultValue = Present(JoinType.Inner))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(type)
+    }
+}
+
+internal class LeftJoinWith : AbstractJoinWith() {
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(JoinType.Left)
+    }
+}
+
+internal class RightJoinWith : AbstractJoinWith() {
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(JoinType.Right)
+    }
+}
+
+internal class FullJoinWith : AbstractJoinWith() {
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(JoinType.Full)
+    }
+}
+
+internal class InnerJoinWith : AbstractJoinWith() {
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(JoinType.Inner)
+    }
+}
+
+internal class FilterJoinWith : AbstractJoinWith() {
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(JoinType.Filter)
+    }
+}
+
+internal class ExcludeJoinWith : AbstractJoinWith() {
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return join(JoinType.Exclude)
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/map.kt

@@ -0,0 +1,18 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dsl
+
+class MapToFrame : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.body by dsl()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val addDsl = AddDslApproximation(mutableListOf())
+        body(addDsl, emptyMap())
+        return PluginDataFrameSchema(addDsl.columns)
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/merge.kt

@@ -0,0 +1,104 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+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.remove
+import org.jetbrains.kotlinx.dataframe.api.replace
+import org.jetbrains.kotlinx.dataframe.api.toPath
+import org.jetbrains.kotlinx.dataframe.api.under
+import org.jetbrains.kotlinx.dataframe.api.with
+import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
+import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
+import org.jetbrains.kotlinx.dataframe.impl.ColumnNameGenerator
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.type
+
+class MergeApproximation(
+    val df: PluginDataFrameSchema,
+    val columns: ColumnsResolver,
+)
+
+class Merge0 : AbstractInterpreter<MergeApproximation>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.selector: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): MergeApproximation {
+        return MergeApproximation(receiver, selector)
+    }
+}
+
+class MergeInto0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: MergeApproximation by arg()
+    val Arguments.columnName: String by arg()
+    val Arguments.typeArg2 by type()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = receiver.columns.resolve(receiver.df).map { it.path.toPath() }
+        return merge(receiver.df, columns, pathOf(columnName), simpleColumnOf(columnName, typeArg2.type))
+    }
+}
+
+class MergeId : AbstractInterpreter<MergeApproximation>() {
+    val Arguments.receiver: MergeApproximation by arg()
+
+    override fun Arguments.interpret(): MergeApproximation {
+        return receiver
+    }
+}
+
+class MergeBy0 : AbstractInterpreter<MergeApproximation>() {
+    val Arguments.receiver: MergeApproximation by arg()
+    val Arguments.separator by ignore()
+    val Arguments.prefix by ignore()
+    val Arguments.postfix by ignore()
+    val Arguments.limit by ignore()
+    val Arguments.truncated by ignore()
+
+    override fun Arguments.interpret(): MergeApproximation {
+        return receiver
+    }
+}
+
+class MergeBy1 : AbstractInterpreter<MergeApproximation>() {
+    val Arguments.receiver: MergeApproximation by arg()
+    val Arguments.infer by ignore()
+    val Arguments.transform by ignore()
+
+    override fun Arguments.interpret(): MergeApproximation {
+        return receiver
+    }
+}
+
+fun merge(
+    schema: PluginDataFrameSchema,
+    columns: List<ColumnPath>,
+    path: ColumnPath,
+    result: SimpleCol
+): PluginDataFrameSchema {
+    val df = schema.asDataFrame()
+    val mergedPath = if (df.getColumnOrNull(path) != null) {
+        val temp = ColumnNameGenerator(df.columnNames()).addUnique("temp")
+        pathOf(temp)
+    } else {
+        path
+    }
+
+    val grouped = df.move { columns.toColumnSet() }.under { mergedPath }
+
+    var res = grouped.replace { mergedPath }.with { result.rename(mergedPath.columnName).asDataColumn() }
+    if (mergedPath != path) {
+        res = res.remove { path }.move { mergedPath }.into { path }
+    }
+    return res.toPluginDataFrameSchema()
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/move.kt

@@ -0,0 +1,118 @@
+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.moveToStart
+import org.jetbrains.kotlinx.dataframe.api.pathOf
+import org.jetbrains.kotlinx.dataframe.api.to
+import org.jetbrains.kotlinx.dataframe.api.toStart
+import org.jetbrains.kotlinx.dataframe.api.toEnd
+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
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+
+class Move0 : AbstractInterpreter<MoveClauseApproximation>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): MoveClauseApproximation {
+        return MoveClauseApproximation(receiver, columns)
+    }
+}
+
+class ToTop : 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() }.toTop().toPluginDataFrameSchema()
+    }
+}
+
+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{ pathOf(column) }.toPluginDataFrameSchema()
+    }
+}
+
+class MoveToStart0 : 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() }.toStart().toPluginDataFrameSchema()
+    }
+}
+
+class MoveToStart1 : 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().moveToStart { columns.toColumnSet() }.toPluginDataFrameSchema()
+    }
+}
+
+
+class MoveToEnd0 : 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() }.toEnd().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 MoveTo : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: MoveClauseApproximation by arg()
+    val Arguments.columnIndex: Int by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.df.asDataFrame().move { receiver.columns }.to(columnIndex).toPluginDataFrameSchema()
+    }
+}
+
+class MoveClauseApproximation(val df: PluginDataFrameSchema, val columns: ColumnsResolver)

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/remove.kt

@@ -0,0 +1,43 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+
+class Remove0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val removeResult = removeImpl(receiver.columns(), columns.resolve(receiver).mapTo(mutableSetOf()) { it.path.path })
+        return PluginDataFrameSchema(removeResult.updatedColumns)
+    }
+}
+
+class RemoveResult(val updatedColumns: List<SimpleCol>, val removedColumns: List<SimpleCol>)
+
+fun KotlinTypeFacade.removeImpl(schema: List<SimpleCol>, paths: Set<List<String>>): RemoveResult {
+    val removed = mutableListOf<SimpleCol>()
+    fun remove(schema: List<SimpleCol>, p: List<String>): List<SimpleCol> {
+        return schema.flatMap {
+            if (p + it.name() in paths) {
+                removed.add(it)
+                emptyList()
+            } else {
+                if (it is SimpleColumnGroup) {
+                    val columns = remove(it.columns(), p + it.name())
+                    if (columns.isEmpty()) emptyList() else listOf(SimpleColumnGroup(it.name(), columns))
+                } else {
+                    listOf(it)
+                }
+            }
+        }
+    }
+    val updatedColumns = remove(schema, emptyList())
+    return RemoveResult(updatedColumns, removed)
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/rename.kt

@@ -0,0 +1,137 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
+import org.jetbrains.kotlinx.dataframe.api.rename
+import org.jetbrains.kotlinx.dataframe.api.renameToCamelCase
+import org.jetbrains.kotlinx.dataframe.api.toCamelCase
+import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+
+class Rename : AbstractInterpreter<RenameClauseApproximation>() {
+    private val Arguments.receiver by dataFrame()
+    private val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): RenameClauseApproximation = RenameClauseApproximation(receiver, columns)
+}
+
+class RenameClauseApproximation(val schema: PluginDataFrameSchema, val columns: ColumnsResolver)
+
+class RenameInto : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: RenameClauseApproximation by arg()
+    val Arguments.newNames: List<String> by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = receiver.columns.resolve(receiver.schema)
+        require(columns.size == newNames.size)
+        var i = 0
+        return receiver.schema.map(
+            selected = columns.mapTo(mutableSetOf()) { it.path.path },
+            nextName = { newNames[i].also { i += 1 } },
+        )
+    }
+}
+
+class RenameMapping : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by dataFrame()
+    val Arguments.mappings: List<Interpreter.Success<Pair<*, *>>> by arg()
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val mappings = mappings.map {
+            val it = it.value
+            val name = (it.first as? FirLiteralExpression)?.value as? String
+            val newName = (it.second as? FirLiteralExpression)?.value  as? String
+            if (name == null || newName == null) return PluginDataFrameSchema(emptyList())
+            name to newName
+        }
+
+        return receiver.asDataFrame().rename(*mappings.toTypedArray()).toPluginDataFrameSchema()
+    }
+}
+
+internal fun PluginDataFrameSchema.map(selected: ColumnsSet, nextName: () -> String): PluginDataFrameSchema =
+    PluginDataFrameSchema(
+        f(columns(), nextName, selected, emptyList()),
+    )
+
+internal fun f(
+    columns: List<SimpleCol>,
+    nextName: () -> String,
+    selected: ColumnsSet,
+    path: List<String>,
+): List<SimpleCol> =
+    columns.map {
+        val fullPath = path + listOf(it.name)
+        when (it) {
+            is SimpleColumnGroup -> {
+                val group = if (fullPath in selected) {
+                    it.rename(nextName())
+                } else {
+                    it
+                }
+                group.map(selected, fullPath, nextName)
+            }
+
+            is SimpleFrameColumn -> {
+                val frame = if (fullPath in selected) {
+                    it.rename(nextName())
+                } else {
+                    it
+                }
+                frame.map(selected, fullPath, nextName)
+            }
+
+            is SimpleDataColumn -> if (fullPath in selected) {
+                it.rename(nextName())
+            } else {
+                it
+            }
+        }
+    }
+
+internal fun SimpleColumnGroup.map(
+    selected: ColumnsSet,
+    path: List<String>,
+    nextName: () -> String,
+): SimpleColumnGroup =
+    SimpleColumnGroup(
+        name = name,
+        columns = f(columns(), nextName, selected, path),
+    )
+
+internal fun SimpleFrameColumn.map(
+    selected: ColumnsSet,
+    path: List<String>,
+    nextName: () -> String,
+): SimpleFrameColumn =
+    SimpleFrameColumn(
+        name = name,
+        columns = f(columns(), nextName, selected, path),
+    )
+
+class RenameToCamelCase : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by dataFrame()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema =
+        receiver.asDataFrame().renameToCamelCase().toPluginDataFrameSchema()
+}
+
+class RenameToCamelCaseClause : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: RenameClauseApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val selectedPaths = receiver.columns.resolve(receiver.schema).map { it.path }
+        return receiver.schema.asDataFrame()
+            .rename { selectedPaths.toColumnSet() }.toCamelCase()
+            .toPluginDataFrameSchema()
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/reorder.kt

@@ -0,0 +1,43 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.api.byName
+import org.jetbrains.kotlinx.dataframe.api.reorder
+import org.jetbrains.kotlinx.dataframe.api.reorderColumnsByName
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+
+class ReorderColumnsByName : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by dataFrame()
+    val Arguments.atAnyDepth: Boolean by arg(defaultValue = Present(true))
+    val Arguments.desc: Boolean by arg(defaultValue = Present(false))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.asDataFrame().reorderColumnsByName(atAnyDepth, desc).toPluginDataFrameSchema()
+    }
+}
+
+class Reorder : AbstractInterpreter<ReorderApproximation>() {
+    val Arguments.receiver by dataFrame()
+    val Arguments.selector: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ReorderApproximation {
+        return ReorderApproximation(receiver, selector)
+    }
+}
+
+class ReorderApproximation(val df: PluginDataFrameSchema, val selector: ColumnsResolver)
+
+class ByName : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: ReorderApproximation by arg()
+    val Arguments.desc: Boolean by arg(defaultValue = Present(false))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.df.asDataFrame().reorder { receiver.selector }.byName(desc).toPluginDataFrameSchema()
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/select.kt

@@ -0,0 +1,715 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.fir.types.isMarkedNullable
+import org.jetbrains.kotlin.fir.types.isSubtypeOf
+import org.jetbrains.kotlinx.dataframe.api.asColumnGroup
+import org.jetbrains.kotlinx.dataframe.impl.columns.ColumnsList
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asSimpleColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
+import org.jetbrains.kotlinx.dataframe.plugin.impl.type
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+
+/**
+ * NOTE: Serves both, select and distinct operations.
+ */
+internal class Select0 : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return PluginDataFrameSchema(columns.resolve(receiver).map { it.column })
+    }
+}
+
+internal class Expr0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.name: String by arg(defaultValue = Present("untitled"))
+    val Arguments.infer by ignore()
+    val Arguments.expression: TypeApproximation by type()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return SingleColumnApproximation(
+            ColumnWithPathApproximation(
+                ColumnPathApproximation(listOf(name)),
+                SimpleDataColumn(name, expression)
+            )
+        )
+    }
+}
+
+internal class And0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.other: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return object : ColumnsResolver, ColumnsList<Any?> {
+            override val columns = listOf(receiver, other)
+
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                return receiver.resolve(df) + other.resolve(df)
+            }
+        }
+    }
+}
+
+internal class All0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.all() }
+    }
+}
+
+internal class All1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { all() }
+    }
+}
+
+internal class All2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.asColumnGroup().allCols()
+        }
+    }
+}
+
+internal class AllAfter0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.column: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.allAfter(column) }
+    }
+}
+
+internal class AllAfter1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.column: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { allAfter(column) }
+    }
+}
+
+internal class AllAfter2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.column: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { allAfter(column) }
+    }
+}
+
+internal class AllAfter3 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.column: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.asColumnGroup().allColsAfter(column)
+        }
+    }
+}
+
+internal class AllBefore0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.column: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.allBefore(column) }
+    }
+}
+
+internal class AllBefore1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.column: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { allBefore(column) }
+    }
+}
+
+internal class AllBefore2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.column: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { allBefore(column) }
+    }
+}
+
+internal class AllUpTo0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.column: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.allUpTo(column) }
+    }
+}
+
+internal class AllUpTo1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.column: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { allUpTo(column) }
+    }
+}
+
+internal class AllUpTo2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.column: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { allUpTo(column) }
+    }
+}
+
+internal class AllFrom0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.column: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.allFrom(column) }
+    }
+}
+
+internal class AllFrom1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.column: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { allFrom(column) }
+    }
+}
+
+internal class AllFrom2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.column: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { allFrom(column) }
+    }
+}
+
+internal class ColsOf0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.typeArg0: TypeApproximation by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return object : ColumnsResolver {
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                val cols = df.columns().map {
+                    val path = ColumnPathApproximation(listOf(it.name))
+                    ColumnWithPathApproximation(path, it)
+                }
+                return colsOf(cols, typeArg0.type)
+            }
+        }
+    }
+
+}
+
+private fun Arguments.colsOf(cols: List<ColumnWithPathApproximation>, type: ConeKotlinType) =
+    cols
+        .filter {
+            it.column.isColOf(type, session)
+        }
+
+fun SimpleCol.isColOf(type: ConeKotlinType, session: FirSession): Boolean {
+    val columnType = when (this) {
+        is SimpleDataColumn -> this.type.type
+        is SimpleColumnGroup -> ConeClassLikeTypeImpl(
+            ConeClassLikeLookupTagImpl(Names.DATA_ROW_CLASS_ID),
+            typeArguments = arrayOf(session.builtinTypes.anyType.type),
+            isNullable = false
+        )
+        is SimpleFrameColumn -> ConeClassLikeTypeImpl(
+            ConeClassLikeLookupTagImpl(Names.DF_CLASS_ID),
+            typeArguments = arrayOf(session.builtinTypes.anyType.type),
+            isNullable = false
+        )
+    }
+    return columnType.isSubtypeOf(type, session)
+}
+
+internal class ColsAtAnyDepth0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.colsAtAnyDepth() }
+    }
+}
+
+internal class ColsAtAnyDepth1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { colsAtAnyDepth() }
+    }
+}
+
+internal class ColsAtAnyDepth2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.asColumnGroup().colsAtAnyDepth()
+        }
+    }
+}
+
+internal class ColsOf1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.typeArg0: TypeApproximation by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return object : ColumnsResolver {
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                return colsOf(receiver.resolve(df), typeArg0.type)
+            }
+        }
+    }
+}
+
+internal class ColsOf2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.typeArg0: TypeApproximation by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.all().filter { it.asSimpleColumn().isColOf(typeArg0.type, session) }
+        }
+    }
+}
+
+internal class WithoutNulls0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return object : ColumnsResolver {
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                return receiver.resolve(df).filter {
+                    withoutNulls(it.column)
+                }
+            }
+        }
+    }
+}
+
+private fun withoutNulls(col: SimpleCol): Boolean = col !is SimpleDataColumn || !col.type.type.isMarkedNullable
+
+internal class WithoutNulls1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return object : ColumnsResolver {
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                val cols = df.columns().map {
+                    val path = ColumnPathApproximation(listOf(it.name))
+                    ColumnWithPathApproximation(path, it)
+                }
+                return cols.filter {
+                    withoutNulls(it.column)
+                }
+            }
+        }
+    }
+}
+
+internal class WithoutNulls2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.filter {
+                withoutNulls(it.data.asSimpleColumn())
+            }
+        }
+    }
+}
+
+internal class FrameCols0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.frameCols() }
+    }
+}
+
+internal class FrameCols1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { frameCols() }
+    }
+}
+
+internal class FrameCols2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.asColumnGroup().frameCols() }
+    }
+}
+
+internal class ColGroups0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.colGroups() }
+    }
+}
+
+internal class ColGroups1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { colGroups() }
+    }
+}
+
+internal class ColGroups2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.asColumnGroup().colGroups() }
+    }
+}
+
+internal class NameContains0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.text: String by arg()
+    val Arguments.ignoreCase: Boolean by arg(defaultValue = Present(false))
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.nameContains(text, ignoreCase) }
+    }
+}
+
+internal class NameContains1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.text: String by arg()
+    val Arguments.ignoreCase: Boolean by arg(defaultValue = Present(false))
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { nameContains(text, ignoreCase) }
+    }
+}
+
+internal class NameContains2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.text: String by arg()
+    val Arguments.ignoreCase: Boolean by arg(defaultValue = Present(false))
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.asColumnGroup().colsNameContains(text, ignoreCase)
+        }
+    }
+}
+
+internal class NameStartsWith0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.prefix: String by arg()
+    val Arguments.ignoreCase: Boolean by arg(defaultValue = Present(false))
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.nameStartsWith(prefix, ignoreCase) }
+    }
+}
+
+internal class NameStartsWith1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.prefix: String by arg()
+    val Arguments.ignoreCase: Boolean by arg(defaultValue = Present(false))
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { nameStartsWith(prefix, ignoreCase) }
+    }
+}
+
+internal class NameStartsWith2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.prefix: String by arg()
+    val Arguments.ignoreCase: Boolean by arg(defaultValue = Present(false))
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.asColumnGroup().colsNameStartsWith(prefix, ignoreCase)
+        }
+    }
+}
+
+internal class NameEndsWith0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.suffix: String by arg()
+    val Arguments.ignoreCase: Boolean by arg(defaultValue = Present(false))
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.nameEndsWith(suffix, ignoreCase) }
+    }
+}
+
+internal class NameEndsWith1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.suffix: String by arg()
+    val Arguments.ignoreCase: Boolean by arg(defaultValue = Present(false))
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { nameEndsWith(suffix, ignoreCase) }
+    }
+}
+
+internal class NameEndsWith2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.suffix: String by arg()
+    val Arguments.ignoreCase: Boolean by arg(defaultValue = Present(false))
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.asColumnGroup().colsNameEndsWith(suffix, ignoreCase)
+        }
+    }
+}
+
+internal class ColumnRange : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.endInclusive: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver..endInclusive
+        }
+    }
+}
+
+internal class Cols0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.firstCol: SingleColumnApproximation by arg()
+    val Arguments.otherCols: List<Interpreter.Success<*>> by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { cols(firstCol, *otherCols.map { it.value as SingleColumnApproximation }.toTypedArray()) }
+    }
+}
+
+internal class Single0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.single() }
+    }
+}
+
+internal class Single1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { single() }
+    }
+}
+
+internal class Single2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.asColumnGroup().singleCol()
+        }
+    }
+}
+
+internal class First0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.first() }
+    }
+}
+
+internal class First1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { first() }
+    }
+}
+
+internal class First2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.asColumnGroup().firstCol()
+        }
+    }
+}
+
+internal class Last0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.last() }
+    }
+}
+
+internal class Last1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { last() }
+    }
+}
+
+internal class Last2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.asColumnGroup().lastCol()
+        }
+    }
+}
+
+internal class Take0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.n: Int by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.take(n) }
+    }
+}
+
+internal class Take1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.n: Int by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { take(n) }
+    }
+}
+
+internal class Take2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.n: Int by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.asColumnGroup().takeCols(n)
+        }
+    }
+}
+
+internal class TakeLast0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.n: Int by arg(defaultValue = Present(1))
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.takeLast(n) }
+    }
+}
+
+internal class TakeLast1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.n: Int by arg(defaultValue = Present(1))
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { takeLast(n) }
+    }
+}
+
+internal class TakeLast2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.n: Int by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.asColumnGroup().takeLastCols(n)
+        }
+    }
+}
+
+internal class Drop0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.n: Int by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.drop(n) }
+    }
+}
+
+internal class Drop1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.n: Int by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { drop(n) }
+    }
+}
+
+internal class Drop2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.n: Int by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.asColumnGroup().dropCols(n)
+        }
+    }
+}
+
+internal class DropLast0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.n: Int by arg(defaultValue = Present(1))
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.dropLast(n) }
+    }
+}
+
+internal class DropLast1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    val Arguments.n: Int by arg(defaultValue = Present(1))
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { dropLast(n) }
+    }
+}
+
+internal class DropLast2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.n: Int by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.asColumnGroup().dropLastCols(n)
+        }
+    }
+}
+
+internal class ValueCols0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.valueCols() }
+    }
+}
+
+internal class ValueCols1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { valueCols() }
+    }
+}
+
+internal class ValueCols2 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver {
+            receiver.valueCols()
+        }
+    }
+}
+
+internal class Named0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.newName: String by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver named newName }
+    }
+}
+
+internal class NestedSelect : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.selector: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.asColumnGroup().select { selector } }
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/statistics.kt

@@ -0,0 +1,172 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeNullability
+import org.jetbrains.kotlin.fir.types.constructClassLikeType
+import org.jetbrains.kotlin.fir.types.isSubtypeOf
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
+import kotlin.reflect.KClass
+import kotlin.reflect.KType
+import kotlin.reflect.full.createType
+
+private object PrimitiveClassIds {
+    const val INT = "kotlin/Int"
+    const val LONG = "kotlin/Long"
+    const val DOUBLE = "kotlin/Double"
+    const val FLOAT = "kotlin/Float"
+    const val SHORT = "kotlin/Short"
+    const val BYTE = "kotlin/Byte"
+}
+
+private fun KClass<*>.toClassId(): ClassId? = when (this) {
+    Int::class -> ClassId.fromString(PrimitiveClassIds.INT)
+    Long::class -> ClassId.fromString(PrimitiveClassIds.LONG)
+    Double::class -> ClassId.fromString(PrimitiveClassIds.DOUBLE)
+    Float::class -> ClassId.fromString(PrimitiveClassIds.FLOAT)
+    Short::class -> ClassId.fromString(PrimitiveClassIds.SHORT)
+    Byte::class -> ClassId.fromString(PrimitiveClassIds.BYTE)
+    else -> null
+}
+
+private val primitiveTypeMap = mapOf(
+    PrimitiveClassIds.INT to Int::class,
+    PrimitiveClassIds.LONG to Long::class,
+    PrimitiveClassIds.DOUBLE to Double::class,
+    PrimitiveClassIds.FLOAT to Float::class,
+    PrimitiveClassIds.SHORT to Short::class,
+    PrimitiveClassIds.BYTE to Byte::class
+)
+
+fun ConeKotlinType.toKType(): KType? {
+    return (this as? ConeClassLikeType)?.let { coneType ->
+        val nullable = coneType.nullability == ConeNullability.NULLABLE
+        primitiveTypeMap[coneType.lookupTag.classId.asString()]
+            ?.createType(nullable = nullable)
+    }
+}
+
+fun KType.toConeKotlinType(): ConeKotlinType? {
+    val kClass = this.classifier as? KClass<*> ?: return null
+    val classId = kClass.toClassId() ?: return null
+
+    return classId.constructClassLikeType(
+        typeArguments = emptyArray(),
+        isNullable = this.isMarkedNullable
+    )
+}
+
+internal fun Arguments.generateStatisticResultColumns(
+    statisticAggregator: Aggregator<*, *>,
+    inputColumns: List<SimpleDataColumn>
+): List<SimpleCol> {
+    return inputColumns.map { col -> createUpdatedColumn(col, statisticAggregator) }
+}
+
+private fun Arguments.createUpdatedColumn(
+    column: SimpleDataColumn,
+    statisticAggregator: Aggregator<*, *>
+): SimpleCol {
+    val originalType = column.type.type
+    val inputKType = originalType.toKType()
+    val resultKType = inputKType?.let { statisticAggregator.calculateReturnType(it, emptyInput = true) }
+    val updatedType = resultKType?.toConeKotlinType() ?: originalType
+    return simpleColumnOf(column.name, updatedType)
+}
+
+internal val skipNaN = true
+internal val ddofDefault: Int = 1
+internal val percentileArg: Double = 30.0
+internal val sum = Aggregators.sum(skipNaN)
+internal val mean = Aggregators.mean(skipNaN)
+internal val std = Aggregators.std(skipNaN, ddofDefault)
+internal val median = Aggregators.median(skipNaN)
+internal val min = Aggregators.min<Double>(skipNaN)
+internal val max = Aggregators.max<Double>(skipNaN)
+internal val percentile = Aggregators.percentile(percentileArg, skipNaN)
+
+/** Adds to the schema only numerical columns. */
+abstract class Aggregator0(val aggregator: Aggregator<*, *>) : AbstractSchemaModificationInterpreter() {
+    private val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val resolvedColumns = receiver.columns()
+            .filterIsInstance<SimpleDataColumn>()
+            .filter { it.type.type.isSubtypeOf(session.builtinTypes.numberType.type, session) }
+
+        val newColumns = generateStatisticResultColumns(aggregator, resolvedColumns)
+
+        return PluginDataFrameSchema(newColumns)
+    }
+}
+
+class Sum0 : Aggregator0(sum)
+
+class Mean0 : Aggregator0(mean)
+
+class Std0 : Aggregator0(std)
+
+/** Adds to the schema only numerical columns. */
+abstract class AggregatorIntraComparable0(val aggregator: Aggregator<*, *>) : AbstractSchemaModificationInterpreter() {
+    private val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val resolvedColumns = receiver.columns()
+            .filterIsInstance<SimpleDataColumn>()
+            .filter { isIntraComparable(it, session) }
+
+        val newColumns = generateStatisticResultColumns(aggregator, resolvedColumns)
+
+        return PluginDataFrameSchema(newColumns)
+    }
+}
+
+class Median0 : AggregatorIntraComparable0(median)
+
+class Max0 : AggregatorIntraComparable0(max)
+
+class Min0 : AggregatorIntraComparable0(min)
+
+class Percentile0 : AggregatorIntraComparable0(percentile) {
+    val Arguments.percentile by ignore()
+}
+
+/** Adds to the schema all resolved columns. */
+abstract class Aggregator1 (val aggregator: Aggregator<*, *>) : AbstractSchemaModificationInterpreter() {
+    private val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    private val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val resolvedColumns = columns.resolve(receiver).map { it.column }.filterIsInstance<SimpleDataColumn>().toList()
+
+        val newColumns = generateStatisticResultColumns(aggregator, resolvedColumns)
+
+        return PluginDataFrameSchema(newColumns)
+    }
+}
+
+class Sum1 : Aggregator1(sum)
+
+class Mean1 : Aggregator1(mean)
+
+class Std1 : Aggregator1(std)
+
+class Median1 : Aggregator1(median)
+
+class Max1 : Aggregator1(max)
+
+class Min1 : Aggregator1(min)
+
+class Percentile1 : Aggregator1(percentile) {
+    val Arguments.percentile by ignore()
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/stdlib.kt

@@ -0,0 +1,40 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+
+class PairToConstructor : AbstractInterpreter<Pair<*, *>>() {
+    val Arguments.receiver: Any? by arg(lens = Interpreter.Id)
+    val Arguments.that: Any? by arg(lens = Interpreter.Id)
+    override fun Arguments.interpret(): Pair<*, *> {
+        return receiver to that
+    }
+}
+
+class PairConstructor : AbstractInterpreter<Pair<*, *>>() {
+    val Arguments.first: Any? by arg(lens = Interpreter.Id)
+    val Arguments.second: Any? by arg(lens = Interpreter.Id)
+    override fun Arguments.interpret(): Pair<*, *> {
+        return first to second
+    }
+}
+
+class TrimMargin : AbstractInterpreter<String>() {
+    val Arguments.receiver: String by arg()
+    val Arguments.marginPrefix: String by arg(defaultValue = Present("|"))
+
+    override fun Arguments.interpret(): String {
+        return receiver.trimMargin(marginPrefix)
+    }
+}
+
+class TrimIndent : AbstractInterpreter<String>() {
+    val Arguments.receiver: String by arg()
+
+    override fun Arguments.interpret(): String {
+        return receiver.trimIndent()
+    }
+}
+

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/toDataFrame.kt

@@ -0,0 +1,410 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlinx.dataframe.plugin.classId
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.hasAnnotation
+import org.jetbrains.kotlin.fir.declarations.utils.isEnumClass
+import org.jetbrains.kotlin.fir.declarations.utils.isStatic
+import org.jetbrains.kotlin.fir.declarations.utils.visibility
+import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
+import org.jetbrains.kotlin.fir.expressions.FirVarargArgumentsExpression
+import org.jetbrains.kotlin.fir.java.JavaTypeParameterStack
+import org.jetbrains.kotlin.fir.java.declarations.FirJavaClass
+import org.jetbrains.kotlin.fir.java.resolveIfJavaType
+import org.jetbrains.kotlin.fir.references.toResolvedPropertySymbol
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.scopes.collectAllFunctions
+import org.jetbrains.kotlin.fir.scopes.collectAllProperties
+import org.jetbrains.kotlin.fir.scopes.unsubstitutedScope
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.ConeFlexibleType
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeNullability
+import org.jetbrains.kotlin.fir.types.ConeStarProjection
+import org.jetbrains.kotlin.fir.types.ConeTypeParameterType
+import org.jetbrains.kotlin.fir.types.ConeTypeProjection
+import org.jetbrains.kotlin.fir.types.canBeNull
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.constructClassLikeType
+import org.jetbrains.kotlin.fir.types.isArrayTypeOrNullableArrayType
+import org.jetbrains.kotlin.fir.types.isNullable
+import org.jetbrains.kotlin.fir.types.isStarProjection
+import org.jetbrains.kotlin.fir.types.isSubtypeOf
+import org.jetbrains.kotlin.fir.types.resolvedType
+import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
+import org.jetbrains.kotlin.fir.types.toSymbol
+import org.jetbrains.kotlin.fir.types.type
+import org.jetbrains.kotlin.fir.types.typeContext
+import org.jetbrains.kotlin.fir.types.upperBoundIfFlexible
+import org.jetbrains.kotlin.fir.types.withArguments
+import org.jetbrains.kotlin.fir.types.withNullability
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.StandardClassIds
+import org.jetbrains.kotlin.name.StandardClassIds.List
+import org.jetbrains.kotlinx.dataframe.codeGen.*
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.wrap
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dsl
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.type
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names.DATA_ROW_CLASS_ID
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names.DATA_SCHEMA_CLASS_ID
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names.DF_CLASS_ID
+import java.util.*
+
+class ToDataFrameDsl : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: FirExpression? by arg(lens = Interpreter.Id)
+    val Arguments.body by dsl()
+    val Arguments.typeArg0: ConeTypeProjection? by arg(lens = Interpreter.Id)
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val dsl = CreateDataFrameDslImplApproximation()
+        body(dsl, mapOf("typeArg0" to Interpreter.Success(typeArg0)))
+        return PluginDataFrameSchema(dsl.columns)
+    }
+}
+
+class ToDataFrame : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: FirExpression? by arg(lens = Interpreter.Id)
+    val Arguments.maxDepth: Number by arg(defaultValue = Present(DEFAULT_MAX_DEPTH))
+    val Arguments.typeArg0: ConeTypeProjection by arg(lens = Interpreter.Id)
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return toDataFrame(maxDepth.toInt(), typeArg0, TraverseConfiguration())
+    }
+}
+
+class ToDataFrameDefault : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: FirExpression? by arg(lens = Interpreter.Id)
+    val Arguments.typeArg0: ConeTypeProjection by arg(lens = Interpreter.Id)
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return toDataFrame(DEFAULT_MAX_DEPTH, typeArg0, TraverseConfiguration())
+    }
+}
+
+class ToDataFrameColumn : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: FirExpression? by arg(lens = Interpreter.Id)
+    val Arguments.typeArg0 by type()
+    val Arguments.columnName: String by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return PluginDataFrameSchema(listOf(simpleColumnOf(columnName, typeArg0.type)))
+    }
+}
+
+private const val DEFAULT_MAX_DEPTH = 0
+
+class Properties0 : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: CreateDataFrameDslImplApproximation by arg()
+    val Arguments.maxDepth: Int by arg()
+    val Arguments.body by dsl()
+    val Arguments.typeArg0: ConeTypeProjection by arg(lens = Interpreter.Id)
+
+    override fun Arguments.interpret() {
+        dsl.configuration.maxDepth = maxDepth
+        body(dsl.configuration.traverseConfiguration, emptyMap())
+        val schema = toDataFrame(dsl.configuration.maxDepth, typeArg0, dsl.configuration.traverseConfiguration)
+        dsl.columns.addAll(schema.columns())
+    }
+}
+
+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()
+}
+
+class TraverseConfiguration {
+    val excludeProperties = mutableSetOf<FirCallableReferenceAccess>()
+    val excludeClasses = mutableSetOf<FirGetClassCall>()
+    val preserveClasses = mutableSetOf<FirGetClassCall>()
+    val preserveProperties = mutableSetOf<FirCallableReferenceAccess>()
+}
+
+class Preserve0 : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: TraverseConfiguration by arg()
+    val Arguments.classes: FirVarargArgumentsExpression by arg(lens = Interpreter.Id)
+
+    override fun Arguments.interpret() {
+        dsl.preserveClasses.addAll(classes.arguments.filterIsInstance<FirGetClassCall>())
+    }
+}
+
+class Preserve1 : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: TraverseConfiguration by arg()
+    val Arguments.properties: FirVarargArgumentsExpression by arg(lens = Interpreter.Id)
+
+    override fun Arguments.interpret() {
+        dsl.preserveProperties.addAll(properties.arguments.filterIsInstance<FirCallableReferenceAccess>())
+    }
+}
+
+class Exclude0 : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: TraverseConfiguration by arg()
+    val Arguments.classes: FirVarargArgumentsExpression by arg(lens = Interpreter.Id)
+
+    override fun Arguments.interpret() {
+        dsl.excludeClasses.addAll(classes.arguments.filterIsInstance<FirGetClassCall>())
+    }
+}
+
+class Exclude1 : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: TraverseConfiguration by arg()
+    val Arguments.properties: FirVarargArgumentsExpression by arg(lens = Interpreter.Id)
+
+    override fun Arguments.interpret() {
+        dsl.excludeProperties.addAll(properties.arguments.filterIsInstance<FirCallableReferenceAccess>())
+    }
+}
+
+@Suppress("INVISIBLE_MEMBER")
+
+@OptIn(SymbolInternals::class)
+internal fun KotlinTypeFacade.toDataFrame(
+    maxDepth: Int,
+    arg: ConeTypeProjection,
+    traverseConfiguration: TraverseConfiguration,
+): PluginDataFrameSchema {
+    val excludes =
+        traverseConfiguration.excludeProperties.mapNotNullTo(mutableSetOf()) { it.calleeReference.toResolvedPropertySymbol() }
+    val excludedClasses = traverseConfiguration.excludeClasses.mapTo(mutableSetOf()) { it.argument.resolvedType }
+    val preserveClasses = traverseConfiguration.preserveClasses.mapNotNullTo(mutableSetOf()) { it.classId }
+    val preserveProperties =
+        traverseConfiguration.preserveProperties.mapNotNullTo(mutableSetOf()) { it.calleeReference.toResolvedPropertySymbol() }
+
+    fun convert(classLike: ConeKotlinType, depth: Int, makeNullable: Boolean): List<SimpleCol> {
+        val symbol = classLike.toRegularClassSymbol(session) ?: return emptyList()
+        val scope = symbol.unsubstitutedScope(session, ScopeSession(), false, FirResolvePhase.STATUS)
+        val declarations = if (symbol.fir is FirJavaClass) {
+            scope
+                .collectAllFunctions()
+                .filter { !it.isStatic && it.valueParameterSymbols.isEmpty() && it.typeParameterSymbols.isEmpty() }
+                .mapNotNull { function ->
+                    val name = function.name.identifier
+                    if (name.startsWith("get") || name.startsWith("is")) {
+                        val propertyName = name
+                            .replaceFirst("get", "")
+                            .replaceFirst("is", "")
+                            .let {
+                                if (it.firstOrNull()?.isUpperCase() == true) {
+                                    it.replaceFirstChar { it.lowercase(Locale.getDefault()) }
+                                } else {
+                                    null
+                                }
+                            }
+                        propertyName?.let { function to it }
+                    } else {
+                        null
+                    }
+                }
+        } else {
+            scope
+                .collectAllProperties()
+                .filterIsInstance<FirPropertySymbol>()
+                .map {
+                    it to it.name.identifier
+                }
+        }
+
+        return declarations
+            .filterNot { excludes.contains(it.first) }
+            .filterNot { excludedClasses.contains(it.first.resolvedReturnType) }
+            .filter { it.first.visibility == Visibilities.Public }
+            .map { (it, name) ->
+                var returnType = it.fir.returnTypeRef.resolveIfJavaType(session, JavaTypeParameterStack.EMPTY, null)
+                    .coneType.upperBoundIfFlexible()
+
+                returnType = if (returnType is ConeTypeParameterType) {
+                    if (returnType.canBeNull(session)) {
+                        session.builtinTypes.nullableAnyType.type
+                    } else {
+                        session.builtinTypes.anyType.type
+                    }
+                } else {
+                    returnType.withArguments {
+                        val type = it.type
+                        if (type is ConeTypeParameterType) {
+                            session.builtinTypes.nullableAnyType.type
+                        } else {
+                            type?.upperBoundIfFlexible() ?: it
+                        }
+                    }
+                }
+
+                val fieldKind = returnType.getFieldKind(session)
+
+                val keepSubtree =
+                    depth >= maxDepth && !fieldKind.shouldBeConvertedToColumnGroup && !fieldKind.shouldBeConvertedToFrameColumn
+                if (keepSubtree || returnType.isValueType(session) || returnType.classId in preserveClasses || it in preserveProperties) {
+                    SimpleDataColumn(
+                        name,
+                        TypeApproximation(
+                            returnType.withNullability(
+                                ConeNullability.create(makeNullable),
+                                session.typeContext
+                            )
+                        )
+                    )
+                } else if (
+                    returnType.isSubtypeOf(
+                        StandardClassIds.Iterable.constructClassLikeType(arrayOf(ConeStarProjection)),
+                        session
+                    ) ||
+                    returnType.isSubtypeOf(
+                        StandardClassIds.Iterable.constructClassLikeType(
+                            arrayOf(ConeStarProjection),
+                            isNullable = true
+                        ), session
+                    )
+                ) {
+                    val type: ConeKotlinType = when (val typeArgument = returnType.typeArguments[0]) {
+                        is ConeKotlinType -> typeArgument
+                        ConeStarProjection -> session.builtinTypes.nullableAnyType.type
+                        else -> session.builtinTypes.nullableAnyType.type
+                    }
+                    if (type.isValueType(session)) {
+                        val columnType = List.constructClassLikeType(arrayOf(type), returnType.isNullable)
+                            .withNullability(ConeNullability.create(makeNullable), session.typeContext)
+                            .wrap()
+                        SimpleDataColumn(name, columnType)
+                    } else {
+                        SimpleFrameColumn(name, convert(type, depth + 1, makeNullable = false))
+                    }
+                } else {
+                    SimpleColumnGroup(name, convert(returnType, depth + 1, returnType.isNullable || makeNullable))
+                }
+            }
+    }
+
+    arg.type?.let { type ->
+        if (!type.canBeUnfolded(session)) {
+            return PluginDataFrameSchema(listOf(simpleColumnOf("value", type)))
+        }
+    }
+
+    return when {
+        arg.isStarProjection -> PluginDataFrameSchema.EMPTY
+        else -> {
+            val classLike = when (val type = arg.type) {
+                is ConeClassLikeType -> type
+                is ConeFlexibleType -> type.upperBound
+                else -> null
+            } ?: return PluginDataFrameSchema.EMPTY
+            val columns = convert(classLike, 0, makeNullable = classLike.isNullable)
+            PluginDataFrameSchema(columns)
+        }
+    }
+}
+
+fun ConeKotlinType.canBeUnfolded(session: FirSession): Boolean =
+    !isValueType(session) && hasProperties(session)
+
+private fun ConeKotlinType.isValueType(session: FirSession) =
+    this.isArrayTypeOrNullableArrayType ||
+        this.classId == StandardClassIds.Unit ||
+        this.classId == StandardClassIds.Any ||
+        this.classId == StandardClassIds.Map ||
+        this.classId == StandardClassIds.MutableMap ||
+        this.classId == StandardClassIds.String ||
+        this.classId in StandardClassIds.primitiveTypes ||
+        this.classId in StandardClassIds.unsignedTypes ||
+        classId in setOf(
+        Names.DURATION_CLASS_ID,
+        Names.LOCAL_DATE_CLASS_ID,
+        Names.LOCAL_DATE_TIME_CLASS_ID,
+        Names.INSTANT_CLASS_ID,
+        Names.DATE_TIME_PERIOD_CLASS_ID,
+        Names.DATE_TIME_UNIT_CLASS_ID,
+        Names.TIME_ZONE_CLASS_ID
+    ) ||
+        this.isSubtypeOf(
+            StandardClassIds.Number.constructClassLikeType(emptyArray(), isNullable = true),
+            session
+        ) ||
+        this.toRegularClassSymbol(session)?.isEnumClass ?: false ||
+        this.isSubtypeOf(
+            Names.TEMPORAL_ACCESSOR_CLASS_ID.constructClassLikeType(emptyArray(), isNullable = true), session
+        ) ||
+        this.isSubtypeOf(
+            Names.TEMPORAL_AMOUNT_CLASS_ID.constructClassLikeType(emptyArray(), isNullable = true), session
+        )
+
+
+private fun ConeKotlinType.hasProperties(session: FirSession): Boolean {
+    val symbol = this.toRegularClassSymbol(session) as? FirClassSymbol<*> ?: return false
+    val scope = symbol.unsubstitutedScope(
+        session,
+        ScopeSession(),
+        withForcedTypeCalculator = false,
+        memberRequiredPhase = null
+    )
+
+    return scope.collectAllProperties().any { it.visibility == Visibilities.Public } ||
+        scope.collectAllFunctions().any { it.visibility == Visibilities.Public && it.isGetterLike() }
+}
+
+private fun FirNamedFunctionSymbol.isGetterLike(): Boolean {
+    val functionName = this.name.asString()
+    return (functionName.startsWith("get") || functionName.startsWith("is")) &&
+        this.valueParameterSymbols.isEmpty() &&
+        this.typeParameterSymbols.isEmpty()
+}
+
+// org.jetbrains.kotlinx.dataframe.codeGen.getFieldKind
+private fun ConeKotlinType.getFieldKind(session: FirSession) = FieldKind.of(
+    this,
+    isDataFrame = { classId == DF_CLASS_ID },
+    isListToFrame = { classId == List && typeArguments[0].type.hasAnnotation(DATA_SCHEMA_CLASS_ID, session) },
+    isDataRow = { classId == DATA_ROW_CLASS_ID },
+    isObjectToGroup = { hasAnnotation(DATA_SCHEMA_CLASS_ID, session) }
+)
+
+private fun ConeKotlinType?.hasAnnotation(id: ClassId, session: FirSession) =
+    this?.toSymbol(session)?.hasAnnotation(id, session) == true
+
+
+class CreateDataFrameDslImplApproximation {
+    val configuration: CreateDataFrameConfiguration = CreateDataFrameConfiguration()
+    val columns: MutableList<SimpleCol> = mutableListOf()
+}
+
+class ToDataFrameFrom : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: CreateDataFrameDslImplApproximation by arg()
+    val Arguments.receiver: String by arg()
+    val Arguments.expression: TypeApproximation by type()
+    override fun Arguments.interpret() {
+        dsl.columns += simpleColumnOf(receiver, expression.type)
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/unfold.kt

@@ -0,0 +1,38 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.api.replace
+import org.jetbrains.kotlinx.dataframe.api.with
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asSimpleColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+
+class DataFrameUnfold : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.properties by ignore()
+    val Arguments.maxDepth: Int by arg(defaultValue = Present(0))
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.asDataFrame().replace { columns }.with {
+            val column = it.asSimpleColumn() as? SimpleDataColumn
+            if (column != null) {
+                if (!column.type.type.canBeUnfolded(session)) {
+                    it
+                } else {
+                    SimpleColumnGroup(it.name(), toDataFrame(maxDepth, column.type.type, TraverseConfiguration()).columns()).asDataColumn()
+                }
+            } else {
+                it
+            }
+        }.toPluginDataFrameSchema()
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/ungroup.kt

@@ -0,0 +1,34 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+
+class Ungroup0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val ungrouped = ungroupImpl(receiver.columns(), columns.resolve(receiver).mapTo(mutableSetOf()) { it.path.path }, emptyList())
+        return PluginDataFrameSchema(ungrouped)
+    }
+}
+
+fun KotlinTypeFacade.ungroupImpl(schema: List<SimpleCol>, path: Set<List<String>>, p: List<String>): List<SimpleCol> {
+    return schema.flatMap {
+        if (it !is SimpleColumnGroup) {
+            listOf(it)
+        } else {
+            if (p + it.name in path) {
+                it.columns()
+            } else {
+                listOf(SimpleColumnGroup(it.name, ungroupImpl(it.columns(), path, p + it.name)))
+            }
+        }
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/update.kt

@@ -0,0 +1,19 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+
+class Update0 : AbstractInterpreter<UpdateApproximationImpl>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): UpdateApproximationImpl {
+        return UpdateApproximationImpl(receiver, columns)
+    }
+}
+
+sealed interface UpdateApproximation
+
+class UpdateApproximationImpl(val schema: PluginDataFrameSchema, val columns: ColumnsResolver) : UpdateApproximation

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/valueCounts.kt

@@ -0,0 +1,28 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.api.ValueCount
+import org.jetbrains.kotlinx.dataframe.impl.ColumnNameGenerator
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.wrap
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
+
+class ValueCounts : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by dataFrame()
+    val Arguments.dropNA by ignore()
+    val Arguments.ascending by ignore()
+    val Arguments.sort by ignore()
+    val Arguments.resultColumn: String by arg(defaultValue = Present(ValueCount::count.name))
+    val Arguments.columns: ColumnsResolver? by arg(defaultValue = Present(null))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val res = columns?.resolve(receiver)?.map { it.column } ?: receiver.columns()
+        val generator = ColumnNameGenerator(res.map { it.name })
+        val count = SimpleDataColumn(generator.addUnique(resultColumn), session.builtinTypes.intType.type.wrap())
+        return PluginDataFrameSchema(res + count)
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/xs.kt

@@ -0,0 +1,52 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirVarargArgumentsExpression
+import org.jetbrains.kotlinx.dataframe.api.getColumnsWithPaths
+import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
+import org.jetbrains.kotlinx.dataframe.api.remove
+import org.jetbrains.kotlinx.dataframe.api.toPath
+import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.groupBy
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+
+class DataFrameXs : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.keyValues: FirExpression by arg(lens = Interpreter.Id)
+    val Arguments.keyColumns: ColumnsResolver? by arg(defaultValue = Present(null))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val keyColumns = keyColumns?.let { it.resolve(receiver).map { it.path.toPath() } }
+        val n = (keyValues as? FirVarargArgumentsExpression)?.arguments?.size ?: return PluginDataFrameSchema.EMPTY
+        return receiver
+            .asDataFrame()
+            .remove { keyColumns?.toColumnSet() ?: colsAtAnyDepth { !it.isColumnGroup() }.take(n) }
+            .toPluginDataFrameSchema()
+    }
+}
+
+class GroupByXs : AbstractInterpreter<GroupBy>() {
+    val Arguments.receiver by groupBy()
+    val Arguments.keyValues: FirExpression by arg(lens = Interpreter.Id)
+    val Arguments.keyColumns: ColumnsResolver? by arg(defaultValue = Present(null))
+
+    override fun Arguments.interpret(): GroupBy {
+        val keyColumns = keyColumns?.let { it.resolve(receiver.keys).map { it.path.toPath() } }
+        val n = (keyValues as? FirVarargArgumentsExpression)?.arguments?.size ?: return GroupBy.EMPTY
+
+        val toRemove = receiver.keys.asDataFrame()
+            .getColumnsWithPaths { keyColumns?.toColumnSet() ?: colsAtAnyDepth { !it.isColumnGroup() }.take(n) }
+            .toColumnSet()
+        val updatedKeys = receiver.keys.asDataFrame().remove { toRemove }.toPluginDataFrameSchema()
+        val updatedGroups = receiver.groups.asDataFrame().remove { toRemove }.toPluginDataFrameSchema()
+        return GroupBy(updatedKeys, updatedGroups)
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/data/ColumnAccessorApproximation.kt

@@ -0,0 +1,5 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.data
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
+
+class ColumnAccessorApproximation(val name: String, val type: TypeApproximation)

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/data/ColumnPathApproximation.kt

@@ -0,0 +1,5 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.data
+
+import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
+
+typealias ColumnPathApproximation = ColumnPath

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/data/ColumnWithPathApproximation.kt

@@ -0,0 +1,9 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.data
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
+
+/**
+ * @see ColumnWithPath
+ */
+data class ColumnWithPathApproximation(val path: ColumnPathApproximation, val column: SimpleCol)

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/data/DataFrameCallableId.kt

@@ -0,0 +1,7 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.data
+
+data class DataFrameCallableId(
+    val packageName: String,
+    val className: String,
+    val callableName: String
+)

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/data/InsertClauseApproximation.kt

@@ -0,0 +1,10 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.data
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.api.InsertClause
+
+/**
+ * @see InsertClause
+*/
+data class InsertClauseApproximation(val df: PluginDataFrameSchema, val column: SimpleCol)

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/data/KPropertyApproximation.kt

@@ -0,0 +1,5 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.data
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
+
+data class KPropertyApproximation(val name: String, val type: TypeApproximation)

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/interpret.kt

@@ -0,0 +1,576 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlinx.dataframe.plugin
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.checkers.toClassLikeSymbol
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.Marker
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.findArgumentByName
+import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
+import org.jetbrains.kotlin.fir.declarations.hasAnnotation
+import org.jetbrains.kotlin.fir.declarations.utils.isLocal
+import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
+import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
+import org.jetbrains.kotlin.fir.expressions.FirErrorExpression
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
+import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
+import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
+import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
+import org.jetbrains.kotlin.fir.expressions.FirThisReceiverExpression
+import org.jetbrains.kotlin.fir.expressions.FirVarargArgumentsExpression
+import org.jetbrains.kotlin.fir.expressions.arguments
+import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentList
+import org.jetbrains.kotlin.fir.references.FirResolvedCallableReference
+import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
+import org.jetbrains.kotlin.fir.references.resolved
+import org.jetbrains.kotlin.fir.references.symbol
+import org.jetbrains.kotlin.fir.references.toResolvedCallableSymbol
+import org.jetbrains.kotlin.fir.resolve.fqName
+import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
+import org.jetbrains.kotlin.fir.scopes.collectAllProperties
+import org.jetbrains.kotlin.fir.scopes.getProperties
+import org.jetbrains.kotlin.fir.scopes.impl.declaredMemberScope
+import org.jetbrains.kotlin.fir.symbols.impl.FirEnumEntrySymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirTypeAliasSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirValueParameterSymbol
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.ConeIntersectionType
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeStarProjection
+import org.jetbrains.kotlin.fir.types.ConeTypeParameterType
+import org.jetbrains.kotlin.fir.types.ConeTypeProjection
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.isNullableString
+import org.jetbrains.kotlin.fir.types.isPrimitiveOrNullablePrimitive
+import org.jetbrains.kotlin.fir.types.isStarProjection
+import org.jetbrains.kotlin.fir.types.isString
+import org.jetbrains.kotlin.fir.types.resolvedType
+import org.jetbrains.kotlin.fir.types.returnType
+import org.jetbrains.kotlin.fir.types.toConeTypeProjection
+import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
+import org.jetbrains.kotlin.fir.types.toSymbol
+import org.jetbrains.kotlin.fir.types.type
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.SessionContext
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.DataFrameCallableId
+import org.jetbrains.kotlinx.dataframe.plugin.impl.data.KPropertyApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColumnsResolver
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBy
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.SingleColumnApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
+
+fun <T> KotlinTypeFacade.interpret(
+    functionCall: FirFunctionCall,
+    processor: Interpreter<T>,
+    additionalArguments: Map<String, Interpreter.Success<Any?>> = emptyMap(),
+    reporter: InterpretationErrorReporter,
+): Interpreter.Success<T>? {
+    val refinedArguments: RefinedArguments = functionCall.collectArgumentExpressions()
+
+    val defaultArguments = processor.expectedArguments.filter { it.defaultValue is Present }.map { it.name }.toSet() + THIS_CALL
+    val actualValueArguments = refinedArguments.associateBy { it.name.identifier }.toSortedMap()
+    val conflictingKeys = additionalArguments.keys intersect actualValueArguments.keys
+    if (conflictingKeys.isNotEmpty()) {
+        if (isTest) {
+            interpretationFrameworkError("Conflicting keys: $conflictingKeys")
+        }
+        return null
+    }
+    val expectedArgsMap = processor.expectedArguments
+        .associateBy { it.name }.toSortedMap().minus(additionalArguments.keys)
+
+    val typeArguments = buildMap {
+        functionCall.typeArguments.forEachIndexed { index, firTypeProjection ->
+            val key = "typeArg$index"
+            val lens = expectedArgsMap[key]?.lens ?: return@forEachIndexed
+            val value: Any = if (lens == Interpreter.Id) {
+                firTypeProjection.toConeTypeProjection()
+            } else {
+                val type = firTypeProjection.toConeTypeProjection().type ?: session.builtinTypes.nullableAnyType.type
+                if (type is ConeIntersectionType) return@forEachIndexed
+                Marker(type)
+            }
+            put(key, Interpreter.Success(value))
+        }
+    }
+
+    val unexpectedArguments = (expectedArgsMap.keys - defaultArguments) != (actualValueArguments.keys + typeArguments.keys - defaultArguments)
+    if (unexpectedArguments) {
+        if (isTest) {
+            val message = buildString {
+                appendLine("ERROR: Different set of arguments")
+                appendLine("Implementation class: $processor")
+                appendLine("Not found in actual: ${expectedArgsMap.keys - actualValueArguments.keys}")
+                val diff = actualValueArguments.keys - expectedArgsMap.keys
+                appendLine("Passed, but not expected: ${diff}")
+                appendLine("add arguments to an interpeter:")
+                appendLine(diff.map { actualValueArguments[it] })
+            }
+            interpretationFrameworkError(message)
+        }
+        return null
+    }
+
+    val arguments = mutableMapOf<String, Interpreter.Success<Any?>>()
+    arguments += additionalArguments
+    arguments += typeArguments
+    arguments[THIS_CALL] = Interpreter.Success(functionCall)
+    val interpretationResults = refinedArguments.refinedArguments.mapNotNull {
+        val name = it.name.identifier
+        val expectedArgument = expectedArgsMap[name] ?: error("$processor $name")
+        val expectedReturnType = expectedArgument.klass
+        val value: Interpreter.Success<Any?>? = when (expectedArgument.lens) {
+            is Interpreter.Value -> {
+                extractValue(it.expression, reporter)
+            }
+
+            is Interpreter.ReturnType -> {
+                val returnType = it.expression.resolvedType.returnType(session)
+                Interpreter.Success(Marker(returnType))
+            }
+
+            is Interpreter.Dsl -> {
+                { receiver: Any, dslArguments: Map<String, Interpreter.Success<Any?>> ->
+                    val map = mapOf("dsl" to Interpreter.Success(receiver)) + dslArguments
+                    (it.expression as FirAnonymousFunctionExpression)
+                        .anonymousFunction.body!!
+                        .statements.filterIsInstance<FirFunctionCall>()
+                        .forEach { call ->
+                            val schemaProcessor = call.loadInterpreter() ?: return@forEach
+                            interpret(
+                                call,
+                                schemaProcessor,
+                                map,
+                                reporter
+                            )
+                        }
+                }.let { Interpreter.Success(it) }
+            }
+
+            is Interpreter.Schema -> {
+                assert(expectedReturnType.toString() == PluginDataFrameSchema::class.qualifiedName!!) {
+                    "'$name' should be ${PluginDataFrameSchema::class.qualifiedName!!}, but plugin expect $expectedReturnType"
+                }
+
+                val objectWithSchema = it.expression.getSchema()
+                if (objectWithSchema == null) {
+                    reporter.doNotReportInterpretationError()
+                    null
+                } else {
+                    val arg = objectWithSchema.schemaArg
+                    val schemaTypeArg = (objectWithSchema.typeRef as ConeClassLikeType).typeArguments[arg]
+                    val schema = pluginDataFrameSchema(schemaTypeArg)
+                    Interpreter.Success(schema)
+                }
+            }
+
+            is Interpreter.GroupBy -> {
+                assert(expectedReturnType.toString() == GroupBy::class.qualifiedName!!) {
+                    "'$name' should be ${GroupBy::class.qualifiedName!!}, but plugin expect $expectedReturnType"
+                }
+                // ok for ReducedGroupBy too
+                val resolvedType = it.expression.resolvedType.fullyExpandedType(session)
+                val keys = pluginDataFrameSchema(resolvedType.typeArguments[0])
+                val groups = pluginDataFrameSchema(resolvedType.typeArguments[1])
+                Interpreter.Success(GroupBy(keys, groups))
+            }
+
+            is Interpreter.Id -> {
+                Interpreter.Success(it.expression)
+            }
+        }
+        value?.let { value1 -> it.name.identifier to value1 }
+    }
+
+    return if (interpretationResults.size == refinedArguments.refinedArguments.size) {
+        arguments.putAll(interpretationResults)
+        when (val res = processor.interpret(arguments, this)) {
+            is Interpreter.Success -> res
+            is Interpreter.Error -> {
+                reporter.reportInterpretationError(functionCall, res.message ?: "")
+                return null
+            }
+        }
+    } else {
+        return null
+    }
+}
+
+private fun KotlinTypeFacade.extractValue(
+    expression: FirExpression?,
+    reporter: InterpretationErrorReporter
+): Interpreter.Success<Any?>? = when (expression) {
+    is FirLiteralExpression -> Interpreter.Success(expression.value!!)
+    is FirVarargArgumentsExpression -> {
+        val args = expression.arguments.map {
+            when (it) {
+                is FirLiteralExpression -> it.value
+                is FirCallableReferenceAccess -> {
+                    toKPropertyApproximation(it, session)
+                }
+
+                is FirFunctionCall -> {
+                    it.loadInterpreter()?.let { processor ->
+                        interpret(it, processor, emptyMap(), reporter)
+                    }
+                }
+
+                else -> extractValue(it, reporter)
+            }
+
+        }
+        Interpreter.Success(args)
+    }
+
+    is FirFunctionCall -> {
+        val interpreter = expression.loadInterpreter()
+        if (interpreter == null) {
+            // if the plugin already transformed call, its original form is the last expression of .let { }
+            val argument = expression.arguments.getOrNull(0)
+            val last = (argument as? FirAnonymousFunctionExpression)?.anonymousFunction?.body?.statements?.lastOrNull()
+            val call = (last as? FirReturnExpression)?.result as? FirFunctionCall
+            call?.loadInterpreter()?.let {
+                interpret(call, it, emptyMap(), reporter)
+            }
+        } else {
+            interpreter.let {
+                val result = interpret(expression, interpreter, emptyMap(), reporter)
+                result
+            }
+        }
+    }
+
+    is FirPropertyAccessExpression -> {
+        (expression.calleeReference as? FirResolvedNamedReference)?.let {
+            val symbol = it.resolvedSymbol
+            val firPropertySymbol = symbol as? FirPropertySymbol
+            val literalInitializer = firPropertySymbol?.resolvedInitializer
+
+            if (symbol is FirEnumEntrySymbol) {
+                Interpreter.Success(
+                    DataFrameCallableId(
+                        packageName = symbol.callableId.packageName.asString(),
+                        className = symbol.callableId.className!!.asString(),
+                        callableName = symbol.callableId.callableName.asString()
+                    )
+                )
+            } else if (literalInitializer != null) {
+                extractValue(literalInitializer, reporter)
+            } else {
+                Interpreter.Success(columnWithPathApproximations(expression))
+            }
+        }
+    }
+
+    is FirCallableReferenceAccess -> {
+        Interpreter.Success(toKPropertyApproximation(expression, session))
+    }
+
+    is FirAnonymousFunctionExpression -> {
+        val result = (expression.anonymousFunction.body?.statements?.lastOrNull() as? FirReturnExpression)?.result
+        val col: Any? = when (result) {
+            is FirPropertyAccessExpression -> {
+                columnWithPathApproximations(result)
+            }
+
+            is FirFunctionCall -> {
+                val interpreter = result.loadInterpreter()
+                if (interpreter == null) {
+                    reporter.reportInterpretationError(result, "Cannot load interpreter")
+                }
+                interpreter?.let {
+                    val value = interpret(result, interpreter, reporter = reporter)?.value
+                    value
+                }
+            }
+
+            is FirErrorExpression -> null
+
+            else -> null
+        }
+        col?.let { Interpreter.Success(it) }
+    }
+
+    else -> null
+}
+
+fun interpretationFrameworkError(message: String): Nothing = throw InterpretationFrameworkError(message)
+
+class InterpretationFrameworkError(message: String) : Error(message)
+
+interface InterpretationErrorReporter {
+    val errorReported: Boolean
+    fun reportInterpretationError(call: FirFunctionCall, message: String)
+
+    fun doNotReportInterpretationError()
+
+    companion object {
+        val DEFAULT = object : InterpretationErrorReporter {
+            override val errorReported: Boolean = false
+            override fun reportInterpretationError(call: FirFunctionCall, message: String) {
+
+            }
+
+            override fun doNotReportInterpretationError() = Unit
+        }
+    }
+}
+
+fun SessionContext.pluginDataFrameSchema(schemaTypeArg: ConeTypeProjection): PluginDataFrameSchema {
+    val schema = if (schemaTypeArg.isStarProjection) {
+        PluginDataFrameSchema.EMPTY
+    } else {
+        val coneClassLikeType = schemaTypeArg.type as? ConeClassLikeType ?: return PluginDataFrameSchema.EMPTY
+        pluginDataFrameSchema(coneClassLikeType)
+    }
+    return schema
+}
+
+fun SessionContext.pluginDataFrameSchema(coneClassLikeType: ConeClassLikeType): PluginDataFrameSchema {
+    val symbol = coneClassLikeType.toSymbol(session) as? FirRegularClassSymbol ?: return PluginDataFrameSchema.EMPTY
+    val declarationSymbols = if (symbol.isLocal && symbol.resolvedSuperTypes.firstOrNull() != session.builtinTypes.anyType.type) {
+        val rootSchemaSymbol = symbol.resolvedSuperTypes.first().toSymbol(session) as? FirRegularClassSymbol
+        rootSchemaSymbol?.declaredMemberScope(session, FirResolvePhase.DECLARATIONS)
+    } else {
+        symbol.declaredMemberScope(session, FirResolvePhase.DECLARATIONS)
+    }.let { scope ->
+        val names = scope?.getCallableNames() ?: emptySet()
+        names.flatMap { scope?.getProperties(it) ?: emptyList() }
+    }
+
+    val mapping = symbol.typeParameterSymbols
+        .mapIndexed { i, symbol -> symbol to coneClassLikeType.typeArguments[i] }
+        .toMap()
+
+    val propertySymbols = declarationSymbols
+        .filterIsInstance<FirPropertySymbol>()
+        .sortPropertiesByOrderAnnotation(sessionContext = this)
+
+    val columns = propertySymbols.mapNotNull { propertySymbol ->
+        columnOf(propertySymbol, mapping)
+    }
+
+    return PluginDataFrameSchema(columns)
+}
+
+private fun List<FirPropertySymbol>.sortPropertiesByOrderAnnotation(sessionContext: SessionContext): List<FirPropertySymbol> {
+    var result = this
+    val annotations = result.mapNotNull {
+        val orderArgument = it.getAnnotationByClassId(
+            Names.ORDER_ANNOTATION,
+            sessionContext.session
+        )?.argumentMapping?.mapping?.get(Names.ORDER_ARGUMENT)
+        (orderArgument as? FirLiteralExpression)?.value as? Int
+    }
+    if (result.size == annotations.size) {
+        result = result.zip(annotations).sortedBy { it.second }.map { it.first }
+    }
+    return result
+}
+
+private fun KotlinTypeFacade.columnWithPathApproximations(result: FirPropertyAccessExpression): ColumnsResolver {
+    return result.resolvedType.let {
+        val column = when (it.classId) {
+            Names.DATA_COLUMN_CLASS_ID -> {
+                val type = when (val arg = it.typeArguments.single()) {
+                    is ConeStarProjection -> session.builtinTypes.nullableAnyType.type
+                    else -> arg as ConeClassLikeType
+                }
+                simpleColumnOf(f(result), type)
+            }
+            Names.COLUM_GROUP_CLASS_ID -> {
+                val arg = it.typeArguments.single()
+                val path = f(result)
+                SimpleColumnGroup(path, pluginDataFrameSchema(arg).columns())
+            }
+            else -> return object : ColumnsResolver {
+                override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                    return emptyList()
+                }
+            }
+        }
+        SingleColumnApproximation(
+            ColumnWithPathApproximation(
+                path = ColumnPathApproximation(path(result)),
+                column
+            )
+        )
+    }
+}
+
+private fun SessionContext.columnOf(it: FirPropertySymbol, mapping: Map<FirTypeParameterSymbol, ConeTypeProjection>): SimpleCol? {
+    val annotation = it.getAnnotationByClassId(Names.COLUMN_NAME_ANNOTATION, session)
+    val columnName = (annotation?.argumentMapping?.mapping?.get(Names.COLUMN_NAME_ARGUMENT) as? FirLiteralExpression)?.value as? String
+    val name = columnName ?: it.name.identifier
+    return when {
+        shouldBeConvertedToFrameColumn(it) -> {
+            val nestedColumns = it.resolvedReturnType.typeArguments[0].type
+                ?.toRegularClassSymbol(session)
+                ?.declaredMemberScope(session, FirResolvePhase.DECLARATIONS)
+                ?.collectAllProperties()
+                ?.filterIsInstance<FirPropertySymbol>()
+                ?.sortPropertiesByOrderAnnotation(this)
+                ?.mapNotNull { columnOf(it, mapping) }
+                ?: emptyList()
+
+            SimpleFrameColumn(name, nestedColumns)
+        }
+
+        shouldBeConvertedToColumnGroup(it) -> {
+            val type = if (isDataRow(it)) it.resolvedReturnType.typeArguments[0].type!! else it.resolvedReturnType
+            val nestedColumns = type
+                .toRegularClassSymbol(session)
+                ?.declaredMemberScope(session, FirResolvePhase.DECLARATIONS)
+                ?.collectAllProperties()
+                ?.filterIsInstance<FirPropertySymbol>()
+                ?.sortPropertiesByOrderAnnotation(this)
+                ?.mapNotNull { columnOf(it, mapping) }
+                ?: emptyList()
+            SimpleColumnGroup(name, nestedColumns)
+        }
+
+        else -> {
+            val type = when (val type = it.resolvedReturnType) {
+                is ConeTypeParameterType -> {
+                    val projection = mapping[type.lookupTag.typeParameterSymbol]
+                    if (projection is ConeStarProjection) {
+                        type.lookupTag.typeParameterSymbol.resolvedBounds.singleOrNull()?.type
+                    } else {
+                        projection as? ConeKotlinType
+                    }
+                }
+
+                else -> type
+            }
+            type?.let { type ->
+                SimpleDataColumn(name, TypeApproximation(type))
+            }
+        }
+    }
+}
+
+private fun SessionContext.shouldBeConvertedToColumnGroup(it: FirPropertySymbol) =
+    isDataRow(it) ||
+        it.resolvedReturnType.toRegularClassSymbol(session)?.hasAnnotation(Names.DATA_SCHEMA_CLASS_ID, session) == true
+
+private fun isDataRow(it: FirPropertySymbol) =
+    it.resolvedReturnType.classId == Names.DATA_ROW_CLASS_ID
+
+private fun SessionContext.shouldBeConvertedToFrameColumn(it: FirPropertySymbol) =
+    isDataFrame(it) ||
+        (it.resolvedReturnType.classId == Names.LIST &&
+            it.resolvedReturnType.typeArguments[0].type?.toRegularClassSymbol(session)?.hasAnnotation(Names.DATA_SCHEMA_CLASS_ID, session) == true)
+
+private fun isDataFrame(it: FirPropertySymbol) =
+    it.resolvedReturnType.classId == Names.DF_CLASS_ID
+
+fun path(propertyAccessExpression: FirPropertyAccessExpression): List<String> {
+    val colName = f(propertyAccessExpression)
+    val typeRef = propertyAccessExpression.dispatchReceiver?.resolvedType
+    val joinDsl = ClassId(FqName("org.jetbrains.kotlinx.dataframe.api"), Name.identifier("JoinDsl"))
+    if (typeRef?.classId?.equals(joinDsl) == true && colName == "right") {
+        return emptyList()
+    }
+    return when (val explicitReceiver = propertyAccessExpression.explicitReceiver) {
+        null, is FirThisReceiverExpression -> listOf(colName)
+        else -> {
+            val propertyAccess = explicitReceiver as FirPropertyAccessExpression
+            if (propertyAccess.calleeReference.symbol is FirValueParameterSymbol) {
+                listOf(colName)
+            } else {
+                path(propertyAccess) + colName
+            }
+        }
+    }
+}
+
+fun f(propertyAccessExpression: FirPropertyAccessExpression): String {
+    return propertyAccessExpression.calleeReference.resolved!!.name.identifier
+}
+
+private fun KotlinTypeFacade.toKPropertyApproximation(
+    firCallableReferenceAccess: FirCallableReferenceAccess,
+    session: FirSession
+): KPropertyApproximation {
+    val propertyName = firCallableReferenceAccess.calleeReference.name.identifier
+    return (firCallableReferenceAccess.calleeReference as FirResolvedCallableReference).let {
+        val symbol = it.toResolvedCallableSymbol()!!
+        val columnName = symbol.annotations
+            .find { it.fqName(session)!!.asString() == ColumnName::class.qualifiedName!! }
+            ?.let {
+                (it.argumentMapping.mapping[Name.identifier(ColumnName::name.name)] as FirLiteralExpression).value as String
+            }
+        val kotlinType = symbol.resolvedReturnTypeRef.type
+
+        val type1 = Marker(kotlinType)
+        KPropertyApproximation(columnName ?: propertyName, type1)
+    }
+}
+
+internal fun FirFunctionCall.collectArgumentExpressions(): RefinedArguments {
+    val refinedArgument = mutableListOf<RefinedArgument>()
+
+    val parameterName = Name.identifier("receiver")
+    (explicitReceiver ?: extensionReceiver)?.let {
+        if (it is FirResolvedQualifier && it.resolvedToCompanionObject) {
+            return@let
+        }
+        refinedArgument += RefinedArgument(parameterName, it)
+    }
+
+    (argumentList as FirResolvedArgumentList).mapping.forEach { (expression, parameter) ->
+        refinedArgument += RefinedArgument(parameter.name, expression)
+    }
+    return RefinedArguments(refinedArgument)
+}
+
+internal val KotlinTypeFacade.getSchema: FirExpression.() -> ObjectWithSchema? get() = { getSchema(session) }
+
+internal fun FirExpression.getSchema(session: FirSession): ObjectWithSchema? {
+    return resolvedType.toSymbol(session)?.let {
+        val (typeRef: ConeKotlinType, symbol) = if (it is FirTypeAliasSymbol) {
+            it.resolvedExpandedTypeRef.coneType to it.resolvedExpandedTypeRef.toClassLikeSymbol(session)!!
+        } else {
+            resolvedType to it
+        }
+        symbol.annotations.firstNotNullOfOrNull {
+            runIf(it.fqName(session)?.asString() == HasSchema::class.qualifiedName!!) {
+                val argumentName = Name.identifier(HasSchema::schemaArg.name)
+                val schemaArg = (it.findArgumentByName(argumentName) as FirLiteralExpression).value
+                ObjectWithSchema((schemaArg as Number).toInt(), typeRef)
+            }
+        } ?: error("Annotate ${symbol} with @HasSchema")
+    }
+}
+
+private const val THIS_CALL = "functionCall"
+
+internal class ObjectWithSchema(val schemaArg: Int, val typeRef: ConeKotlinType)
+
+internal val ConeKotlinType.canHaveLiteralInitializer get() = isPrimitiveOrNullablePrimitive || isString || isNullableString

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt

@@ -0,0 +1,537 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlinx.dataframe.plugin
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names.INTERPRETABLE_FQNAME
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Add
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddWithDsl
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.And0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.And10
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Convert0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Convert2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Convert6
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameGroupBy
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DropNulls0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Exclude0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Exclude1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Explode0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Expr0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.From
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Group0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AggregateDslInto
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByToDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Insert0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Insert1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Into
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Into0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Join0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Match0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Preserve0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Preserve1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Properties0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Remove0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Rename
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameInto
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Select0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.To0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Under0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Under1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Under4
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Ungroup0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.With0
+import org.jetbrains.kotlin.fir.declarations.findArgumentByName
+import org.jetbrains.kotlin.fir.expressions.FirClassReferenceExpression
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
+import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
+import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
+import org.jetbrains.kotlin.fir.expressions.UnresolvedExpressionTypeAccess
+import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
+import org.jetbrains.kotlin.fir.references.toResolvedFunctionSymbol
+import org.jetbrains.kotlin.fir.resolve.fqName
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.coneType
+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
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AggregateRow
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllAfter0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllAfter1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllAfter2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllAfter3
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllBefore0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllBefore1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllBefore2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllFrom0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllFrom1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllFrom2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllUpTo0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllUpTo1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllUpTo2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AsGroupBy
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AsGroupByDefault
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ByName
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColGroups0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColGroups1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColGroups2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Cols0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsAtAnyDepth0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsAtAnyDepth1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsAtAnyDepth2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColumnOfPairs
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColumnRange
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ConcatWithKeys
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ConvertAsColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameBuilderInvoke0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameOf0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameOf3
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameOfPairs
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameUnfold
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameXs
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Drop0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Drop1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Drop2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DropLast0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DropLast1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DropLast2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DropNa0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ExcludeJoin
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ExcludeJoinWith
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FillNulls0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FilterJoin
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FilterJoinWith
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.First0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.First1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.First2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Flatten0
+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.FrameCols1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FullJoin
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FullJoinWith
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByAdd
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByCount0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByInto
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByMax0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByMax1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByMaxOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByMean0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByMean1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByMeanOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByMedian0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByMedian1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByMedianOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByMin0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByMin1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByMinOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByReduceExpression
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByReduceInto
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByReducePredicate
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByXs
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.InnerJoin
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.InsertAfter0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.LeftJoin
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByStd0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByStd1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByStdOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBySum0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBySum1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBySumOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.InnerJoinWith
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.InsertAt
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.JoinWith
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.LeftJoinWith
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MapToFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Max0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Max1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Mean0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Mean1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Median0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Median1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Merge0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeId
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeBy0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeBy1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeInto0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Min0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Min1
+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.MoveTo
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToStart0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToStart1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToEnd0
+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.NameContains0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NameContains1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NameContains2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NameEndsWith0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NameEndsWith1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NameEndsWith2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NameStartsWith0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NameStartsWith1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NameStartsWith2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Named0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NestedSelect
+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.PerRowCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Percentile0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Percentile1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameMapping
+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
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Update0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.UpdateWith0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ValueCounts
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameToCamelCase
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameToCamelCaseClause
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Reorder
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ReorderColumnsByName
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RightJoin
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RightJoinWith
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Single0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Single1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Single2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Std0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Std1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Sum0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Sum1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ValueCols2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Take0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Take1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Take2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TakeLast0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TakeLast1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TakeLast2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToSpecificType
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToSpecificTypePattern
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToSpecificTypeZone
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ValueCols0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ValueCols1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.WithoutNulls0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.WithoutNulls1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.WithoutNulls2
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+
+internal fun FirFunctionCall.loadInterpreter(session: FirSession): Interpreter<*>? {
+    val interpreter = Stdlib.interpreter(this)
+    if (interpreter != null) return interpreter
+    val symbol =
+        (calleeReference as? FirResolvedNamedReference)?.resolvedSymbol as? FirCallableSymbol ?: return null
+    val argName = Name.identifier("interpreter")
+    return symbol.annotations
+        .find { it.fqName(session)?.equals(INTERPRETABLE_FQNAME) ?: false }
+        ?.let { annotation ->
+            val name = (annotation.findArgumentByName(argName) as FirLiteralExpression).value as String
+            name.load<Interpreter<*>>()
+        }
+}
+
+private object Stdlib {
+    private val map: MutableMap<Key, Interpreter<*>> = mutableMapOf()
+    init {
+        register(Names.TO, Names.PAIR, PairToConstructor())
+        register(Names.PAIR_CONSTRUCTOR, Names.PAIR, PairConstructor())
+        register(Names.TRIM_MARGIN, StandardClassIds.String, TrimMargin())
+        register(Names.TRIM_INDENT, StandardClassIds.String, TrimMargin())
+    }
+
+    @OptIn(UnresolvedExpressionTypeAccess::class)
+    fun interpreter(call: FirFunctionCall): Interpreter<*>? {
+        val id = call.calleeReference.toResolvedFunctionSymbol()?.callableId ?: return null
+        val returnType = call.coneTypeOrNull?.classId ?: return null
+        return map[Key(id, returnType)]
+    }
+
+    fun register(id: CallableId, returnType: ClassId, interpreter: Interpreter<*>) {
+        map[Key(id, returnType)] = interpreter
+    }
+}
+
+private data class Key(
+    val id: CallableId,
+    val returnType: ClassId,
+)
+
+internal fun FirFunctionCall.interpreterName(session: FirSession): String? {
+    val symbol =
+        (calleeReference as? FirResolvedNamedReference)?.resolvedSymbol as? FirCallableSymbol ?: return null
+    val argName = Name.identifier("interpreter")
+    return symbol.annotations
+        .find { it.fqName(session)?.equals(INTERPRETABLE_FQNAME) ?: false }
+        ?.let { annotation ->
+            val name = (annotation.findArgumentByName(argName) as FirLiteralExpression).value as String
+            name
+        }
+}
+
+internal val KotlinTypeFacade.loadInterpreter: FirFunctionCall.() -> Interpreter<*>? get() = { this.loadInterpreter(session) }
+
+internal val FirGetClassCall.classId: ClassId?
+    get() {
+        return when (val argument = argument) {
+            is FirResolvedQualifier -> argument.classId!!
+            is FirClassReferenceExpression -> argument.classTypeRef.coneType.classId
+            else -> null
+        }
+    }
+
+internal inline fun <reified T> ClassId.load(): T {
+    val constructor = Class.forName(asFqNameString())
+        .constructors
+        .firstOrNull { constructor -> constructor.parameterCount == 0 }
+        ?: error("Interpreter $this must have an empty constructor")
+
+    return constructor.newInstance() as T
+}
+
+internal inline fun <reified T> String.load(): T {
+    return when (this) {
+        "Add" -> Add()
+        "From" -> From()
+        "Into" -> Into()
+        "AddWithDsl" -> AddWithDsl()
+        "And10" -> And10()
+        "Convert0" -> Convert0()
+        "Convert2" -> Convert2()
+        "Convert6" -> Convert6()
+        "To0" -> To0()
+        "ToSpecificType" -> ToSpecificType()
+        "ToSpecificTypeZone" -> ToSpecificTypeZone()
+        "ToSpecificTypePattern" -> ToSpecificTypePattern()
+        "With0" -> With0()
+        "ConvertAsColumn" -> ConvertAsColumn()
+        "PerRowCol" -> PerRowCol()
+        "Explode0" -> Explode0()
+        "Insert0" -> Insert0()
+        "Insert1" -> Insert1()
+        "Under0" -> Under0()
+        "Under1" -> Under1()
+        "Under4" -> Under4()
+        "InsertAfter0" -> InsertAfter0()
+        "InsertAt" -> InsertAt()
+        "Join0" -> Join0()
+        "LeftJoin" -> LeftJoin()
+        "RightJoin" -> RightJoin()
+        "FullJoin" -> FullJoin()
+        "InnerJoin" -> InnerJoin()
+        "ExcludeJoin" -> ExcludeJoin()
+        "FilterJoin" -> FilterJoin()
+        "JoinWith" -> JoinWith()
+        "LeftJoinWith" -> LeftJoinWith()
+        "RightJoinWith" -> RightJoinWith()
+        "FullJoinWith" -> FullJoinWith()
+        "InnerJoinWith" -> InnerJoinWith()
+        "ExcludeJoinWith" -> ExcludeJoinWith()
+        "FilterJoinWith" -> FilterJoinWith()
+        "Match0" -> Match0()
+        "Rename" -> Rename()
+        "RenameMapping" -> RenameMapping()
+        "Select0" -> Select0()
+        "Distinct0" -> Select0()
+        "NestedSelect" -> NestedSelect()
+        "Expr0" -> Expr0()
+        "And0" -> And0()
+        "Remove0" -> Remove0()
+        "Group0" -> Group0()
+        "Into0" -> Into0()
+        "Ungroup0" -> Ungroup0()
+        "DropNulls0" -> DropNulls0()
+        "DropNa0" -> DropNa0()
+        "Properties0" -> Properties0()
+        "Preserve0" -> Preserve0()
+        "Preserve1" -> Preserve1()
+        "Exclude0" -> Exclude0()
+        "Exclude1" -> Exclude1()
+        "RenameInto" -> RenameInto()
+        "DataFrameGroupBy" -> DataFrameGroupBy()
+        "AsGroupBy" -> AsGroupBy()
+        "AsGroupByDefault" -> AsGroupByDefault()
+        "AggregateDslInto" -> AggregateDslInto()
+        "GroupByToDataFrame" -> GroupByToDataFrame()
+        "GroupByInto" -> GroupByInto()
+        "ToDataFrameFrom0" -> ToDataFrameFrom()
+        "All0" -> All0()
+        "All1" -> All1()
+        "All2" -> All2()
+        "Cols0" -> Cols0()
+        "AllAfter0" -> AllAfter0()
+        "AllAfter1" -> AllAfter1()
+        "AllAfter2" -> AllAfter2()
+        "AllAfter3" -> AllAfter3()
+        "AllBefore0" -> AllBefore0()
+        "AllBefore1" -> AllBefore1()
+        "AllBefore2" -> AllBefore2()
+        "AllUpTo0" -> AllUpTo0()
+        "AllUpTo1" -> AllUpTo1()
+        "AllUpTo2" -> AllUpTo2()
+        "AllFrom0" -> AllFrom0()
+        "AllFrom1" -> AllFrom1()
+        "AllFrom2" -> AllFrom2()
+        "ColsOf0" -> ColsOf0()
+        "ColsOf1" -> ColsOf1()
+        "ColsOf2" -> ColsOf2()
+        "ColsAtAnyDepth0" -> ColsAtAnyDepth0()
+        "ColsAtAnyDepth1" -> ColsAtAnyDepth1()
+        "ColsAtAnyDepth2" -> ColsAtAnyDepth2()
+        "FrameCols0" -> FrameCols0()
+        "FrameCols1" -> FrameCols1()
+        "FrameCols2" -> FrameCols2()
+        "ColGroups0" -> ColGroups0()
+        "ColGroups1" -> ColGroups1()
+        "ColGroups2" -> ColGroups2()
+        "NameContains0" -> NameContains0()
+        "NameContains1" -> NameContains1()
+        "NameContains2" -> NameContains2()
+        "NameStartsWith0" -> NameStartsWith0()
+        "NameStartsWith1" -> NameStartsWith1()
+        "NameStartsWith2" -> NameStartsWith2()
+        "NameEndsWith0" -> NameEndsWith0()
+        "NameEndsWith" -> NameEndsWith1()
+        "NameEndsWith2" -> NameEndsWith2()
+        "First0" -> First0()
+        "First1" -> First1()
+        "First2" -> First2()
+        "Single0" -> Single0()
+        "Single1" -> Single1()
+        "Single2" -> Single2()
+        "Last0" -> Last0()
+        "Last1" -> Last1()
+        "Last2" -> Last2()
+        "Take0" -> Take0()
+        "Take1" -> Take1()
+        "Take2" -> Take2()
+        "TakeLast0" -> TakeLast0()
+        "TakeLast1" -> TakeLast1()
+        "TakeLast2" -> TakeLast2()
+        "Drop0" -> Drop0()
+        "Drop1" -> Drop1()
+        "Drop2" -> Drop2()
+        "DropLast0" -> DropLast0()
+        "DropLast1" -> DropLast1()
+        "DropLast2" -> DropLast2()
+        "WithoutNulls0" -> WithoutNulls0()
+        "WithoutNulls1" -> WithoutNulls1()
+        "WithoutNulls2" -> WithoutNulls2()
+        "ValueCols0" -> ValueCols0()
+        "ValueCols1" -> ValueCols1()
+        "ValueCols2" -> ValueCols2()
+        "ColumnRange" -> ColumnRange()
+        "Named0" -> Named0()
+        "toDataFrameDsl" -> ToDataFrameDsl()
+        "toDataFrame" -> ToDataFrame()
+        "toDataFrameDefault" -> ToDataFrameDefault()
+        "ToDataFrameDslStringInvoke" -> ToDataFrameDslStringInvoke()
+        "DataFrameOf0" -> DataFrameOf0()
+        "DataFrameOfPairs" -> DataFrameOfPairs()
+        "ColumnOfPairs" -> ColumnOfPairs()
+        "DataFrameBuilderInvoke0" -> DataFrameBuilderInvoke0()
+        "ToDataFrameColumn" -> ToDataFrameColumn()
+        "FillNulls0" -> FillNulls0()
+        "UpdateWith0" -> UpdateWith0()
+        "Flatten0" -> Flatten0()
+        "FlattenDefault" -> FlattenDefault()
+        "AddId" -> AddId()
+        "AddDslStringInvoke" -> AddDslStringInvoke()
+        "AddDslNamedGroup" -> AddDslNamedGroup()
+        "MapToFrame" -> MapToFrame()
+        "Move0" -> Move0()
+        "ToTop" -> ToTop()
+        "Update0" -> Update0()
+        "Aggregate" -> Aggregate()
+        "AggregateRow" -> AggregateRow()
+        "DataFrameOf3" -> DataFrameOf3()
+        "ValueCounts" -> ValueCounts()
+        "RenameToCamelCase" -> RenameToCamelCase()
+        "RenameToCamelCaseClause" -> RenameToCamelCaseClause()
+        "MoveUnder0" -> MoveUnder0()
+        "MoveUnder1" -> MoveUnder1()
+        "MoveInto0" -> MoveInto0()
+        "MoveToStart0" -> MoveToStart0()
+        "MoveToStart1" -> MoveToStart1()
+        "MoveToEnd0" -> MoveToEnd0()
+        "MoveAfter0" -> MoveAfter0()
+        "MoveTo" -> MoveTo()
+        "GroupByAdd" -> GroupByAdd()
+        "Merge0" -> Merge0()
+        "MergeInto0" -> MergeInto0()
+        "MergeId" -> MergeId()
+        "MergeBy0" -> MergeBy0()
+        "MergeBy1" -> MergeBy1()
+        "ReorderColumnsByName" -> ReorderColumnsByName()
+        "Reorder" -> Reorder()
+        "ByName" -> ByName()
+        "Sum0" -> Sum0()
+        "Sum1" -> Sum1()
+        "Mean0" -> Mean0()
+        "Mean1" -> Mean1()
+        "Std0" -> Std0()
+        "Std1" -> Std1()
+        "Median0" -> Median0()
+        "Median1" -> Median1()
+        "Min0" -> Min0()
+        "Min1" -> Min1()
+        "Max0" -> Max0()
+        "Max1" -> Max1()
+        "Percentile0" -> Percentile0()
+        "Percentile1" -> Percentile1()
+        "GroupByCount0" -> GroupByCount0()
+        "GroupByMean0" -> GroupByMean0()
+        "GroupByMean1" -> GroupByMean1()
+        "GroupByMeanOf" -> GroupByMeanOf()
+        "GroupByMedian0" -> GroupByMedian0()
+        "GroupByMedian1" -> GroupByMedian1()
+        "GroupByMedianOf" -> GroupByMedianOf()
+        "GroupBySumOf" -> GroupBySumOf()
+        "GroupBySum0" -> GroupBySum0()
+        "GroupBySum1" -> GroupBySum1()
+        "GroupByReducePredicate" -> GroupByReducePredicate()
+        "GroupByReduceExpression" -> GroupByReduceExpression()
+        "GroupByReduceInto" -> GroupByReduceInto()
+        "GroupByMax0" -> GroupByMax0()
+        "GroupByMax1" -> GroupByMax1()
+        "GroupByMaxOf" -> GroupByMaxOf()
+        "GroupByMin0" -> GroupByMin0()
+        "GroupByMin1" -> GroupByMin1()
+        "GroupByMinOf" -> GroupByMinOf()
+        "GroupByStd0" -> GroupByStd0()
+        "GroupByStd1" -> GroupByStd1()
+        "GroupByStdOf" -> GroupByStdOf()
+        "DataFrameXs" -> DataFrameXs()
+        "GroupByXs" -> GroupByXs()
+        "ConcatWithKeys" -> ConcatWithKeys()
+        "DataFrameUnfold" -> DataFrameUnfold()
+        else -> error("$this")
+    } as T
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/Names.kt

@@ -0,0 +1,103 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlinx.dataframe.plugin.utils
+
+import org.jetbrains.kotlin.builtins.StandardNames
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.checkers.fullyExpandedClassId
+import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeStarProjection
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.fir.types.isSubtypeOf
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlinx.dataframe.annotations.ColumnName
+import org.jetbrains.kotlinx.dataframe.annotations.Converter
+import org.jetbrains.kotlinx.dataframe.annotations.Order
+import org.jetbrains.kotlinx.dataframe.annotations.ScopeProperty
+import kotlin.reflect.KClass
+
+object Names {
+    val DF_CLASS_ID: ClassId
+        get() = ClassId.topLevel(FqName.fromSegments(listOf("org", "jetbrains", "kotlinx", "dataframe", "DataFrame")))
+    val GROUP_BY_CLASS_ID: ClassId
+        get() = ClassId.topLevel(FqName.fromSegments(listOf("org", "jetbrains", "kotlinx", "dataframe", "api",  "GroupBy")))
+
+    val COLUM_GROUP_CLASS_ID: ClassId
+        get() = ClassId(FqName("org.jetbrains.kotlinx.dataframe.columns"), Name.identifier("ColumnGroup"))
+    val FRAME_COLUMN_CLASS_ID: ClassId
+        get() = ClassId(FqName("org.jetbrains.kotlinx.dataframe.columns"), Name.identifier("FrameColumn"))
+    val DATA_COLUMN_CLASS_ID: ClassId
+        get() = ClassId(
+            FqName.fromSegments(listOf("org", "jetbrains", "kotlinx", "dataframe")),
+            Name.identifier("DataColumn")
+        )
+    val BASE_COLUMN_CLASS_ID: ClassId
+        get() = ClassId(
+            FqName.fromSegments(listOf("org", "jetbrains", "kotlinx", "dataframe", "columns")),
+            Name.identifier("BaseColumn")
+        )
+    val COLUMNS_CONTAINER_CLASS_ID: ClassId
+        get() = ClassId(
+            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
+        get() = Name.identifier("org.jetbrains.kotlinx.dataframe.annotations")
+    val INTERPRETABLE_FQNAME: FqName
+        get() = FqName("org.jetbrains.kotlinx.dataframe.annotations.Interpretable")
+    private val annotationsPackage = FqName("org.jetbrains.kotlinx.dataframe.annotations")
+    val ORDER_ANNOTATION = ClassId(annotationsPackage, Name.identifier(Order::class.simpleName!!))
+    val CONVERTER_ANNOTATION = ClassId(annotationsPackage, Name.identifier(Converter::class.simpleName!!))
+    val ORDER_ARGUMENT = Name.identifier(Order::order.name)
+    val SCOPE_PROPERTY_ANNOTATION = ClassId(annotationsPackage, Name.identifier(ScopeProperty::class.simpleName!!))
+    val COLUMN_NAME_ANNOTATION = ClassId(annotationsPackage, Name.identifier(ColumnName::class.simpleName!!))
+    val COLUMN_NAME_ARGUMENT = Name.identifier(ColumnName::name.name)
+
+    val DATA_SCHEMA_CLASS_ID = ClassId(annotationsPackage, Name.identifier("DataSchema"))
+    val LIST = ClassId(FqName("kotlin.collections"), Name.identifier("List"))
+    val DURATION_CLASS_ID = kotlin.time.Duration::class.classId()
+    val LOCAL_DATE_CLASS_ID = ClassId(FqName("kotlinx.datetime"), Name.identifier("LocalDate"))
+    val LOCAL_DATE_TIME_CLASS_ID = ClassId(FqName("kotlinx.datetime"), Name.identifier("LocalDateTime"))
+    val INSTANT_CLASS_ID = ClassId(FqName("kotlinx.datetime"), Name.identifier("Instant"))
+    val DATE_TIME_PERIOD_CLASS_ID = ClassId(FqName("kotlinx.datetime"), Name.identifier("DateTimePeriod"))
+    val DATE_TIME_UNIT_CLASS_ID = ClassId(FqName("kotlinx.datetime"), Name.identifier("DateTimeUnit"))
+    val TIME_ZONE_CLASS_ID = ClassId(FqName("kotlinx.datetime"), Name.identifier("TimeZone"))
+    val TEMPORAL_ACCESSOR_CLASS_ID =  ClassId(FqName("java.time.temporal"), Name.identifier("TemporalAccessor"))
+    val TEMPORAL_AMOUNT_CLASS_ID =  ClassId(FqName("java.time.temporal"), Name.identifier("TemporalAmount"))
+
+
+    val PAIR = ClassId(FqName("kotlin"), Name.identifier("Pair"))
+    val PAIR_CONSTRUCTOR = CallableId(FqName("kotlin"), FqName("Pair"), Name.identifier("Pair"))
+    val TO = CallableId(FqName("kotlin"), Name.identifier("to"))
+    val TRIM_MARGIN = CallableId(StandardNames.TEXT_PACKAGE_FQ_NAME, Name.identifier("trimMargin"))
+    val TRIM_INDENT = CallableId(StandardNames.TEXT_PACKAGE_FQ_NAME, Name.identifier("trimIndent"))
+}
+
+private fun KClass<*>.classId(): ClassId {
+    val fqName = this.qualifiedName ?: throw IllegalStateException("KClass does not have a qualified name")
+    val packageFqName = fqName.substringBeforeLast(".", missingDelimiterValue = "")
+    val className = fqName.substringAfterLast(".")
+    return ClassId(FqName(packageFqName), Name.identifier(className))
+}
+
+fun ConeKotlinType.isDataFrame(session: FirSession) =
+    isSubtypeOf(ConeClassLikeTypeImpl(ConeClassLikeLookupTagImpl(Names.DF_CLASS_ID), arrayOf(ConeStarProjection), isNullable = false), session)
+
+fun ConeKotlinType.isGroupBy(session: FirSession) = fullyExpandedClassId(session) == Names.GROUP_BY_CLASS_ID
+
+fun ConeKotlinType.isDataRow(session: FirSession) = fullyExpandedClassId(session) == Names.DATA_ROW_CLASS_ID

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/firFactories.kt

@@ -0,0 +1,92 @@
+package org.jetbrains.kotlinx.dataframe.plugin.utils
+
+import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.descriptors.EffectiveVisibility
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.builder.buildProperty
+import org.jetbrains.kotlin.fir.declarations.builder.buildPropertyAccessor
+import org.jetbrains.kotlin.fir.declarations.builder.buildReceiverParameter
+import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
+import org.jetbrains.kotlin.fir.moduleData
+import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
+import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLookupTagWithFixedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertyAccessorSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.toFirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.DataFramePlugin
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.PropertyName
+
+internal fun FirDeclarationGenerationExtension.generateExtensionProperty(
+    callableId: CallableId,
+    receiverType: ConeClassLikeTypeImpl,
+    propertyName: PropertyName,
+    returnTypeRef: FirResolvedTypeRef,
+    symbol: FirClassSymbol<*>? = null,
+    effectiveVisibility: EffectiveVisibility = EffectiveVisibility.Public,
+    source: KtSourceElement?
+): FirProperty {
+    val firPropertySymbol = FirPropertySymbol(callableId)
+    return buildProperty {
+        this.source = source
+        propertyName.columnNameAnnotation?.let {
+            annotations += it
+        }
+        moduleData = session.moduleData
+        resolvePhase = FirResolvePhase.BODY_RESOLVE
+        origin = FirDeclarationOrigin.Plugin(DataFramePlugin)
+        status = FirResolvedDeclarationStatusImpl(
+                Visibilities.Public,
+                Modality.FINAL,
+                effectiveVisibility
+        )
+        this.returnTypeRef = returnTypeRef
+        receiverParameter = buildReceiverParameter {
+            typeRef = receiverType.toFirResolvedTypeRef()
+        }
+        val classId = callableId.classId
+        if (classId != null) {
+            dispatchReceiverType = if (symbol != null) {
+                ConeClassLikeTypeImpl(
+                    ConeClassLookupTagWithFixedSymbol(classId, symbol),
+                    emptyArray(),
+                    false
+                )
+            } else {
+                ConeClassLikeTypeImpl(
+                    ConeClassLikeLookupTagImpl(classId),
+                    emptyArray(),
+                    false
+                )
+            }
+        }
+        val firPropertyAccessorSymbol = FirPropertyAccessorSymbol()
+        getter = buildPropertyAccessor {
+            moduleData = session.moduleData
+            resolvePhase = FirResolvePhase.BODY_RESOLVE
+            origin = FirDeclarationOrigin.Plugin(DataFramePlugin)
+            this.returnTypeRef = returnTypeRef
+            dispatchReceiverType = receiverType
+            this.symbol = firPropertyAccessorSymbol
+            propertySymbol = firPropertySymbol
+            isGetter = true
+            status = FirResolvedDeclarationStatusImpl(
+                    Visibilities.Public,
+                    Modality.FINAL,
+                    effectiveVisibility
+            )
+        }
+        name = propertyName.identifier
+        this.symbol = firPropertySymbol
+        isVar = false
+        isLocal = false
+    }
+}

dataframe/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/projectOverDataColumnType.kt

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlinx.dataframe.plugin.utils
+
+import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.fir.types.toTypeProjection
+import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
+
+fun FirResolvedTypeRef.projectOverDataColumnType() =
+    ConeClassLikeTypeImpl(
+        ConeClassLikeLookupTagImpl(Names.DATA_COLUMN_CLASS_ID),
+        arrayOf(coneType.toTypeProjection(Variance.INVARIANT)),
+        isNullable = false
+    )

dataframe/plugins/kotlin-dataframe/testData/box/OuterClass.kt

@@ -0,0 +1,33 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+class OuterClass
+
+@org.jetbrains.kotlinx.dataframe.annotations.DataSchema(isOpen = false)
+interface Hello {
+    val name: String
+    val `test name`: InnerClass
+    val nullableProperty: Int?
+    val a: () -> Unit
+    val d: List<List<*>>
+
+    class InnerClass
+}
+
+val ColumnsContainer<Hello>.col1: DataColumn<String> get() = name
+val ColumnsContainer<Hello>.col2: DataColumn<Hello.InnerClass> get() = `test name`
+val ColumnsContainer<Hello>.col3: DataColumn<Int?> get() = nullableProperty
+val ColumnsContainer<Hello>.col4: DataColumn<() -> Unit> get() = a
+val ColumnsContainer<Hello>.col5: DataColumn<List<List<*>>> get() = d
+
+val DataRow<Hello>.row1: String get() = name
+val DataRow<Hello>.row2: Hello.InnerClass get() = `test name`
+val DataRow<Hello>.row3: Int? get() = nullableProperty
+val DataRow<Hello>.row4: () -> Unit get() = a
+val DataRow<Hello>.row5: List<List<*>> get() = d
+
+fun box(): String {
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/Schema.kt

@@ -0,0 +1,26 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+@DataSchema
+interface Schema {
+    val a: Int
+}
+
+fun box(): String {
+    val res = dataFrameOf("a")(1)
+        .cast<Schema>()
+        .add("wwffffwwehirbwerffwffwffwfffffwfffwfwfwfaw") { 42 }
+
+    res.wwffffwwehirbwerffwffwffwfffffwfffwfwfwfaw.print()
+    res.a.print()
+
+    val b = res.convert { a }.with { it.toString() }
+    
+    b.a
+
+//    val res1 = res.conv
+    //res.filter { it }
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/addDsl.kt

@@ -0,0 +1,20 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df = dataFrameOf("a")(1).add {
+        "id" from { it }
+        "group" {
+            "a" from { it }
+        }
+        group("group1") {
+            "b" from { it }
+        }
+    }
+
+    df.group.a
+    df.group1.b
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/addId.kt

@@ -0,0 +1,11 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df = dataFrameOf("a")(1).addId()
+    val i: DataColumn<Int> = df.id
+    val i1: DataColumn<Int> = dataFrameOf("a")(1).addId("i").i
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/aggregateDataFrame.kt

@@ -0,0 +1,14 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val row = dataFrameOf("a" to List(10) { it }).aggregate {
+        maxOf { a } into "max"
+        minOf { a } into "min"
+    }
+    val i: Int = row.max
+    val i1: Int = row.min
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/asGroupBy.kt

@@ -0,0 +1,15 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df = dataFrameOf("i", "group")(1, dataFrameOf("a", "b")(111, 222))
+    val aggregated1 = df.asGroupBy { group }.aggregate { maxOf { a } into "max" }
+    val aggregated2 = df.asGroupBy().aggregate { maxOf { a } into "max" }
+
+    val i: Int = aggregated1.max[0]
+
+    compareSchemas(aggregated1, aggregated2)
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/castTo.kt

@@ -0,0 +1,21 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+import java.net.URI
+
+fun box(): String {
+    val sample = dataFrameOf("full_name", "html_url", "stargazers_count", "topics", "watchers")(
+        "JetBrains/JPS", URI("https://github.com/JetBrains/JPS").toURL(), 23, "[]", 23
+    )
+
+    val organizations = listOf("https://raw.githubusercontent.com/Kotlin/dataframe/master/data/jetbrains_repositories.csv")
+    organizations.forEach { organization ->
+        val df = DataFrame.readCSV(organization).castTo(sample)
+        println(organizations)
+        println("Repositories: ${df.count()}")
+        println("Top 10:")
+        df.sortBy { stargazers_count.desc() }.take(10).print()
+    }
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/castTo_function.kt

@@ -0,0 +1,27 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+import java.io.File
+
+private fun convert(data: List<String>) = data.map { it.split(":") }.toDataFrame {
+    "part1" from { it[0] }
+    "part2" from { it[1].toInt() }
+    "part3" from { it[2] }
+}
+
+fun serialize(data: List<String>, destination: File) {
+    convert(data).writeJson(destination)
+}
+
+fun deserializeAndUse(file: File) {
+    val df = DataFrame.readJson(file).castTo(schemaFrom = ::convert)
+    df.part1.print()
+}
+
+fun box(): String {
+    val file = File.createTempFile("temp", "json")
+    serialize(listOf("b:1:abc", "c:2:bca"), file)
+    deserializeAndUse(file)
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/colKinds/add.kt

@@ -0,0 +1,38 @@
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.*
+import kotlin.reflect.*
+import kotlin.reflect.full.*
+
+inline fun <reified T> runtimeSchema(row: DataRow<T>) = (typeOf<T>().classifier as KClass<*>).memberProperties.associateBy { it.name }
+inline fun <reified T> runtimeSchema(row: DataFrame<T>) = (typeOf<T>().classifier as KClass<*>).memberProperties.associateBy { it.name }
+fun <T> Map<String, T>.col(s: String) = get(s)!!
+
+@DataSchema
+class Record(val string: String)
+
+fun box(): String {
+    val df = dataFrameOf(Record("abc"))
+    val df1 = df.add("row") {
+        val row = dataFrameOf(Record("")).first()
+        row.takeIf { index() % 2 == 0 }
+    }
+
+    require(runtimeSchema(df1.row[0]).col("string").returnType == typeOf<String?>())
+
+    val row = dataFrameOf(Record(""))
+        .add("int") { 1 }
+        .add("double") { 3.0 }
+        .add("char") { 'c' }
+        .group { int and double }.into("g")
+        .group { g and char }.into("f")
+        .first()
+
+    val df2 = df.add("row") {
+        row.takeIf { index() % 2 == 0 }
+    }
+
+    require(runtimeSchema(df2.row.f.g[0]).col("int").returnType == typeOf<Int?>())
+    require(runtimeSchema(df2.row.f[0]).col("char").returnType == typeOf<Char?>())
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/colKinds/toDataFrame.kt

@@ -0,0 +1,19 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+
+@DataSchema
+data class Record(val str: String)
+
+class A
+
+fun box(): String {
+    val df = dataFrameOf(Record("abc"))
+    val df1 = listOf(A()).toDataFrame {
+        "dataRow" from { df.first() }
+        "dataFrame" from { df }
+    }
+    val a: String = df1.dataRow.str[0]
+    val b: String = df1.dataFrame[0].str[0]
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/columnGroupApi.kt

@@ -0,0 +1,17 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+@DataSchema
+data class Record(val a: String, val b: Int)
+
+fun box(): String {
+    val df = List(10) { Record(it.toString(), it) }.let { dataFrameOf(*it.toTypedArray()) }
+    val group = df.group { a and b }.into("c")
+    val df1 = group.c.add("d") { 1 }
+    df1.a
+    df1.b
+    df1.d
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/columnName.kt

@@ -0,0 +1,19 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+@DataSchema
+data class Record(
+    @ColumnName("a")
+    val abc: String,
+)
+
+fun box(): String {
+    val df = dataFrameOf("a")("1").cast<Record>()
+    df.abc
+
+    val df1 = df.add("b") { 1 }
+    df1.a
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/columnName_invalidSymbol.kt

@@ -0,0 +1,10 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df = dataFrameOf("a.b")(1)
+    df.`a b`
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/columnOf_nested.kt

@@ -0,0 +1,15 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val group = columnOf(
+        "c" to columnOf("2"),
+        "d" to columnOf(123),
+    )
+    val str: DataColumn<String> = group.c
+    val i: DataColumn<Int> = group.d
+
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/columnWithStarProjection.kt

@@ -0,0 +1,16 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+@DataSchema
+data class Record(val id: String, val b: Int)
+
+public data class NameValuePair<V>(val name: String, val value: V)
+
+fun box(): String {
+    val df = dataFrameOf(Record("1", 1), Record("2", 123), Record("3", 321))
+    val df1 = df.first().transpose().dropNulls { value }
+    val v: Any = df1.value[0]
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/concatWithKeys.kt

@@ -0,0 +1,19 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df = dataFrameOf(
+        "value" to listOf(1, 2, 3, 3),
+        "type" to listOf("a", "b", "a", "b")
+    )
+    val gb = df.groupBy { expr { "Category: ${type.uppercase()}" } named "category" }
+    val categoryKey = gb.keys.category
+
+    val dfWithCategory = gb.concatWithKeys()
+
+    val category: DataColumn<String> = dfWithCategory.category
+
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/conflictingJvmDeclarations.kt

@@ -0,0 +1,15 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+@DataSchema
+interface Repo {
+    val name: String
+    val url: String
+}
+//
+@DataSchema
+data class Type(val name: String, val vararg: Boolean)
+
+fun box() = "OK"

dataframe/plugins/kotlin-dataframe/testData/box/convertAsColumn.kt

@@ -0,0 +1,14 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val res = dataFrameOf("a")(1, 2, 3).convert { a }.asColumn { it.convertToString() }
+    val str: DataColumn<String> = res.a
+
+    val res1 = dataFrameOf("a")(1).convert { a }.asColumn { dataFrameOf("b", "c")(2, 3.0).asColumnGroup() }
+    val i: DataColumn<Int> = res1.a.b
+    val d: DataColumn<Double> = res1.a.c
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/convertToDataFrame.kt

@@ -0,0 +1,30 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+@DataSchema
+data class Sessions(
+    val roomId: List<Int>
+)
+
+@DataSchema
+data class Rooms(
+    val id: Int,
+    val sort: Int,
+    val name: String,
+)
+
+class Aaa(val a: List<Rooms>)
+
+fun box(): String {
+    val rooms = dataFrameOf(Rooms(1, 2, "n"))
+    val sessions = dataFrameOf(Sessions(listOf(1, 2)))
+
+    val df = sessions.convert { roomId }.with {
+        listOf(Aaa(listOf(Rooms(1, 2, "n")))).toDataFrame(maxDepth = 2)
+    }
+
+    df.roomId[0].a[0].id
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/convertToShortcuts.kt

@@ -0,0 +1,21 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+import kotlinx.datetime.*
+
+fun box(): String {
+    val daysToStandardMillis = 24 * 60 * 60 * 1000L * 366
+    val df = dataFrameOf("a")(60L * 1000L + daysToStandardMillis).convert { a }.toLocalDateTime(TimeZone.UTC)
+    val localDateTime: LocalDateTime = df.a[0]
+
+    val df1 = dataFrameOf("a")(60L * 1000L + daysToStandardMillis, null).convert { a }.toLocalDateTime(TimeZone.UTC)
+    val localDateTime1: LocalDateTime? = df1.a[0]
+
+    val df2 = dataFrameOf("a")(123).convert { a }.toStr()
+    val str: String = df2.a[0]
+
+    val df3 = dataFrameOf("a")(123, null).convert { a }.toStr()
+    df3.compareSchemas(strict = true)
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/convert_to.kt

@@ -0,0 +1,10 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df = dataFrameOf("col")("1", "2").convert { col }.to<Int>()
+    val i: Int = df.col[0]
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/csDsl/all.kt

@@ -0,0 +1,56 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    compareSchemas(
+        df.select { name and age },
+        df.select { allBefore { city } },
+        df.select { allBefore(city) },
+    )
+
+    compareSchemas(
+        df.select { name and age },
+        df.select { allUpTo { age } },
+        df.select { allUpTo(age) },
+    )
+
+    compareSchemas(
+        df.select { weight and isHappy },
+        df.select { allAfter { city } },
+        df.select { allAfter(city) },
+    )
+
+    compareSchemas(
+        df.select { weight and isHappy },
+        df.select { allFrom { weight } },
+        df.select { allFrom(weight) },
+    )
+
+    compareSchemas(
+        df.select { name and age and city and weight and isHappy },
+        df.select { all() }
+    )
+
+    compareSchemas(
+        df.select { name.firstName and name.lastName },
+        df.select { name.allCols() },
+        df.select { name.allCols().all() },
+    )
+
+    compareSchemas(
+        df.select { weight and isHappy },
+        df.select { all().allAfter(city) },
+//        df.select { all().allAfter { city } },
+    )
+
+    compareSchemas(
+        df.select { age },
+        df.select { colsOf<Int?>().allBefore(weight) },
+        df.select { allBefore(weight).colsOf<Int?>() },
+    )
+
+
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/csDsl/colGroups.kt

@@ -0,0 +1,18 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    compareSchemas(
+        df.select { name },
+        df.select { colGroups() },
+        df.select { nameContains("e").colGroups() },
+    )
+
+    compareSchemas(
+        dfGroup.select { name.firstName },
+        dfGroup.select { name.colsNameContains("Name").colGroups() },
+    )
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/csDsl/cols.kt

@@ -0,0 +1,13 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    compareSchemas(
+        df.select { cols(name, age, city, weight, isHappy) },
+        df.select { all() },
+    )
+
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/csDsl/colsAtAnyDepth.kt

@@ -0,0 +1,20 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df1 = df.select { name }
+
+    compareSchemas(
+        df1.select { name.firstName and name.lastName },
+        df1.select { name.colsAtAnyDepth() },
+    )
+
+    compareSchemas(
+        df1.select { name and name.firstName and name.lastName },
+        df1.select { colsAtAnyDepth() },
+        df1.select { nameStartsWith("name").colsAtAnyDepth() },
+    )
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/csDsl/colsOf.kt

@@ -0,0 +1,29 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    compareSchemas(
+        df.select { city },
+        df.select { colsOf<String?>() },
+        df.select { all().colsOf<String?>() },
+    )
+
+    compareSchemas(
+        df.select { name.firstName and name.lastName },
+        df.select { name.colsOf<String>() },
+    )
+
+    compareSchemas(
+        df.select { name },
+        df.select { colsOf<AnyRow>() }
+    )
+
+    val df1 = dataFrameOf("nestedDf")(dataFrameOf("a", "b")(1, 2))
+    compareSchemas(
+        df1.select { nestedDf },
+        df1.select { colsOf<AnyFrame>() }
+    )
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/csDsl/columnNameFIlters.kt

@@ -0,0 +1,39 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    compareSchemas(
+        df.select { age },
+        df.select { nameContains("age") },
+        df.select { nameContains("AGE", ignoreCase = true) },
+        df.select { all().nameContains("age") },
+    )
+
+    compareSchemas(
+        df.select { name.firstName },
+        df.select { name.colsNameContains("first") },
+        df.select { name.colsNameContains("FIRST", ignoreCase = true) },
+    )
+
+    compareSchemas(
+        df.select { age },
+        df.select { nameStartsWith("age") },
+        df.select { nameStartsWith("AGE", ignoreCase = true) },
+        df.select { all().nameStartsWith("age") },
+    )
+
+    compareSchemas(
+        df.select { name.firstName },
+        df.select { name.colsNameStartsWith("first") },
+        df.select { name.colsNameStartsWith("FIRST", ignoreCase = true) },
+    )
+
+    compareSchemas(
+        df.select { name.firstName and name.lastName },
+        df.select { name.colsNameEndsWith("Name") },
+        df.select { name.colsNameEndsWith("NAME", ignoreCase = true) },
+    )
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/csDsl/columnRange.kt

@@ -0,0 +1,17 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    compareSchemas(
+        df.select { age and city and weight },
+        df.select { age..weight },
+    )
+
+    compareSchemas(
+        df.select { name.firstName and name.lastName },
+        df.select { name.firstName..name.lastName }
+    )
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/csDsl/drop.kt

@@ -0,0 +1,27 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    compareSchemas(
+        df.select { name.lastName },
+        df.select { name.dropCols(1) },
+    )
+
+    compareSchemas(
+        df.select { name.firstName },
+        df.select { name.dropLastCols(1) },
+    )
+
+    compareSchemas(
+        df.select { weight and isHappy },
+        df.select { drop(3) },
+    )
+
+    compareSchemas(
+        df.select { name and age },
+        df.select { dropLast(3) },
+    )
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/csDsl/first.kt

@@ -0,0 +1,18 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    compareSchemas(
+        df.select { name },
+        df.select { first() },
+        df.select { nameStartsWith("name").first() },
+    )
+
+    compareSchemas(
+        df.select { name.firstName },
+        df.select { name.firstCol() },
+    )
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/csDsl/frameCols.kt

@@ -0,0 +1,21 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df1 = dataFrameOf("a", "b", "frameCol")(1, 2, dataFrameOf("e", "f")(3, 4))
+
+    compareSchemas(
+        df1.select { frameCol },
+        df1.select { frameCols() },
+        df1.select { nameContains("a").frameCols() },
+    )
+
+    val into = df1.group { a and frameCol }.into("c")
+    compareSchemas(
+        into.select { c.frameCol },
+        into.select { c.frameCols() },
+    )
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/csDsl/last.kt

@@ -0,0 +1,18 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    compareSchemas(
+        df.select { isHappy },
+        df.select { last() },
+        df.select { nameStartsWith("is").last() },
+    )
+
+    compareSchemas(
+        df.select { name.lastName },
+        df.select { name.lastCol() },
+    )
+    return "OK"
+}

dataframe/plugins/kotlin-dataframe/testData/box/csDsl/rename.kt

@@ -0,0 +1,16 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df1 = df.select { expr { age } into "age2" }
+    val i1: Int = df1.age2[0]
+
+    val df2 = dataFrameOf("lists")(listOf(1, 2), listOf(3)).explode { lists into "int" }
+    val i2: Int = df2.int[0]
+
+    df.select { age named "age2" }.compareSchemas(strict = true)
+
+    return "OK"
+}