commonMain.space.kscience.dataforge.actions.ReduceAction.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of dataforge-data Show documentation
dataforge-data
There is a newer version: 0.7.0
package space.kscience.dataforge.actions

import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.fold
import space.kscience.dataforge.data.*
import space.kscience.dataforge.meta.Meta
import space.kscience.dataforge.meta.MutableMeta
import space.kscience.dataforge.misc.DFBuilder
import space.kscience.dataforge.misc.DFExperimental
import space.kscience.dataforge.misc.DFInternal
import space.kscience.dataforge.names.Name
import kotlin.reflect.KType
import kotlin.reflect.typeOf


public class JoinGroup(public var name: String, internal val set: DataSet) {

    public var meta: MutableMeta = MutableMeta()

    public lateinit var result: suspend ActionEnv.(Map) -> R

    public fun result(f: suspend ActionEnv.(Map) -> R) {
        this.result = f;
    }

}

@DFBuilder
public class ReduceGroupBuilder(
    private val inputType: KType,
    private val scope: CoroutineScope,
    public val actionMeta: Meta,
) {
    private val groupRules: MutableList) -> List>> = ArrayList();

    /**
     * introduce grouping by meta value
     */
    public fun byValue(tag: String, defaultTag: String = "@default", action: JoinGroup.() -> Unit) {
        groupRules += { node ->
            GroupRule.byMetaValue(scope, tag, defaultTag).gather(node).map {
                JoinGroup(it.key, it.value).apply(action)
            }
        }
    }

    public fun group(
        groupName: String,
        filter: suspend (Name, Data) -> Boolean,
        action: JoinGroup.() -> Unit,
    ) {
        groupRules += { source ->
            listOf(
                JoinGroup(groupName, source.filter(filter)).apply(action)
            )
        }
    }

    /**
     * Apply transformation to the whole node
     */
    public fun result(resultName: String, f: suspend ActionEnv.(Map) -> R) {
        groupRules += { node ->
            listOf(JoinGroup(resultName, node).apply { result(f) })
        }
    }

    internal suspend fun buildGroups(input: DataSet): List> {
        return groupRules.flatMap { it.invoke(input) }
    }

}

@PublishedApi
internal class ReduceAction(
    private val inputType: KType,
    outputType: KType,
    private val action: ReduceGroupBuilder.() -> Unit,
) : CachingAction(outputType) {
    //TODO optimize reduction. Currently the whole action recalculates on push


    override fun CoroutineScope.transform(set: DataSet, meta: Meta, key: Name): Flow> = flow {
        ReduceGroupBuilder(inputType, this@transform, meta).apply(action).buildGroups(set).forEach { group ->
            val dataFlow: Map> = group.set.flow().fold(HashMap()) { acc, value ->
                acc.apply {
                    acc[value.name] = value.data
                }
            }

            val groupName: String = group.name

            val groupMeta = group.meta

            val env = ActionEnv(Name.parse(groupName), groupMeta, meta)
            @OptIn(DFInternal::class) val res: Data = dataFlow.reduceToData(
                outputType,
                meta = groupMeta
            ) { group.result.invoke(env, it) }

            emit(res.named(env.name))
        }
    }
}

/**
 * A one-to-one mapping action
 */
@DFExperimental
@Suppress("FunctionName")
public inline fun  Action.Companion.reduce(
    noinline builder: ReduceGroupBuilder.() -> Unit,
): Action = ReduceAction(typeOf(), typeOf(), builder)