diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt index ccb73d6e2..48ca42009 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt @@ -269,6 +269,8 @@ public fun DataFrame.Companion.readJsonStr( * @param header Optional list of column names. If given, [text] will be read like an object with [header] being the keys. * @return [DataRow] from the given [text]. */ +@Refine +@Interpretable("DataRowReadJsonStr") public fun DataRow.Companion.readJsonStr( @Language("json") text: String, header: List = emptyList(), diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt index c7b6af676..e882cfcfd 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt @@ -106,7 +106,7 @@ class FunctionCallTransformer( fun transformOrNull(call: FirFunctionCall, originalSymbol: FirNamedFunctionSymbol): FirFunctionCall? } - private val transformers = listOf(GroupByCallTransformer(), DataFrameCallTransformer()) + private val transformers = listOf(GroupByCallTransformer(), DataFrameCallTransformer(), DataRowCallTransformer()) override fun intercept(callInfo: CallInfo, symbol: FirNamedFunctionSymbol): CallReturnType? { val callSiteAnnotations = (callInfo.callSite as? FirAnnotationContainer)?.annotations ?: emptyList() @@ -156,14 +156,14 @@ class FunctionCallTransformer( ?: call } - inner class DataFrameCallTransformer : CallTransformer { + inner class DataSchemaLikeCallTransformer(val classId: ClassId) : CallTransformer { override fun interceptOrNull(callInfo: CallInfo, symbol: FirNamedFunctionSymbol, hash: String): CallReturnType? { - if (symbol.resolvedReturnType.fullyExpandedClassId(session) != Names.DF_CLASS_ID) return null - // possibly null if explicit receiver type is AnyFrame + 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(Names.DF_CLASS_ID) + val lookupTag = ConeClassLikeLookupTagImpl(classId) val typeRef = buildResolvedTypeRef { type = ConeClassLikeTypeImpl( lookupTag, @@ -182,7 +182,7 @@ class FunctionCallTransformer( @OptIn(SymbolInternals::class) override fun transformOrNull(call: FirFunctionCall, originalSymbol: FirNamedFunctionSymbol): FirFunctionCall? { - val callResult = analyzeRefinedCallShape(call, Names.DF_CLASS_ID, InterpretationErrorReporter.DEFAULT) + val callResult = analyzeRefinedCallShape(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!! @@ -195,6 +195,10 @@ class FunctionCallTransformer( } } + inner class DataFrameCallTransformer : CallTransformer by DataSchemaLikeCallTransformer(Names.DF_CLASS_ID) + + inner class DataRowCallTransformer : CallTransformer by DataSchemaLikeCallTransformer(Names.DATA_ROW_CLASS_ID) + inner class GroupByCallTransformer : CallTransformer { override fun interceptOrNull( callInfo: CallInfo, diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ReturnTypeBasedReceiverInjector.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ReturnTypeBasedReceiverInjector.kt index f3478b1a9..b522951c3 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ReturnTypeBasedReceiverInjector.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ReturnTypeBasedReceiverInjector.kt @@ -19,7 +19,7 @@ class ReturnTypeBasedReceiverInjector(session: FirSession) : FirExpressionResolu @OptIn(SymbolInternals::class) override fun addNewImplicitReceivers(functionCall: FirFunctionCall): List { val callReturnType = functionCall.resolvedType - return if (callReturnType.classId in setOf(Names.DF_CLASS_ID, Names.GROUP_BY_CLASS_ID)) { + return if (callReturnType.classId in setOf(Names.DF_CLASS_ID, Names.GROUP_BY_CLASS_ID, Names.DATA_ROW_CLASS_ID)) { val typeArguments = callReturnType.typeArguments typeArguments .mapNotNull { diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/read.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/read.kt index 282fe2a5f..7a70e7208 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/read.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/read.kt @@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.dataframe.plugin.impl.api import kotlinx.serialization.decodeFromString import kotlinx.serialization.json.Json import org.jetbrains.kotlinx.dataframe.DataFrame +import org.jetbrains.kotlinx.dataframe.DataRow import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments import org.jetbrains.kotlinx.dataframe.plugin.impl.Present @@ -118,6 +119,15 @@ internal class ReadJsonStr : AbstractInterpreter() { } } +internal class DataRowReadJsonStr : AbstractInterpreter() { + val Arguments.text: String by arg() + val Arguments.typeClashTactic: JSON.TypeClashTactic by arg(defaultValue = Present(ARRAY_AND_VALUE_COLUMNS)) + + override fun Arguments.interpret(): PluginDataFrameSchema { + return DataRow.readJsonStr(text, typeClashTactic = typeClashTactic).schema().toPluginDataFrameSchema() + } +} + internal class ReadExcel : AbstractSchemaModificationInterpreter() { val Arguments.fileOrUrl: String by arg() val Arguments.sheetName: String? by arg(defaultValue = Present(null)) diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt index becc8fa92..07871bbd0 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt @@ -81,6 +81,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf1 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.DataRowReadJsonStr import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FillNulls0 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Flatten0 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FlattenDefault @@ -218,6 +219,7 @@ internal inline fun String.load(): T { "DataFrameGroupBy" -> DataFrameGroupBy() "GroupByInto" -> GroupByInto() "ReadJsonStr" -> ReadJsonStr() + "DataRowReadJsonStr" -> DataRowReadJsonStr() "ReadDelimStr" -> ReadDelimStr() "GroupByToDataFrame" -> GroupByToDataFrame() "ToDataFrameFrom0" -> ToDataFrameFrom() diff --git a/plugins/kotlin-dataframe/testData/box/readJsonStr_datarow.kt b/plugins/kotlin-dataframe/testData/box/readJsonStr_datarow.kt new file mode 100644 index 000000000..90902081f --- /dev/null +++ b/plugins/kotlin-dataframe/testData/box/readJsonStr_datarow.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.* + +const val text = """{"a":"abc", "b":1}""" + +fun box(): String { + val row = DataRow.readJsonStr(text) + row.a + row.b + return "OK" +} diff --git a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java index 9d35da9f2..221b3d7a1 100644 --- a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java +++ b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java @@ -352,6 +352,12 @@ public void testReadJsonStr_const() { runTest("testData/box/readJsonStr_const.kt"); } + @Test + @TestMetadata("readJsonStr_datarow.kt") + public void testReadJsonStr_datarow() { + runTest("testData/box/readJsonStr_datarow.kt"); + } + @Test @TestMetadata("readJsonStr_localProperty.kt") public void testReadJsonStr_localProperty() {