groupBy impl for JMESPath (proof-of-concept)

jmespath-groupBy.kt

import com.fasterxml.jackson.databind.JsonNode
import com.fasterxml.jackson.databind.node.ArrayNode
import com.fasterxml.jackson.databind.node.JsonNodeFactory
import com.fasterxml.jackson.databind.node.ObjectNode
import io.burt.jmespath.Adapter
import io.burt.jmespath.JmesPath
import io.burt.jmespath.JmesPathType
import io.burt.jmespath.RuntimeConfiguration
import io.burt.jmespath.function.ArgumentConstraints
import io.burt.jmespath.function.BaseFunction
import io.burt.jmespath.function.FunctionArgument
import io.burt.jmespath.function.FunctionRegistry
import io.burt.jmespath.jackson.JacksonRuntime


fun main() {

    // There's a default registry that contains the built in JMESPath functions
    val defaultFunctions = FunctionRegistry.defaultRegistry()
// And we can create a new registry with additional functions by extending it
    val customFunctions = defaultFunctions.extend(GroupByFunction())
// To configure the runtime with the registry we need to create a configuration
    val configuration = RuntimeConfiguration.Builder()
        .withFunctionRegistry(customFunctions)
        .withSilentTypeErrors(true)
        .build()
// And then create a runtime with the configuration
    val runtime: JmesPath<JsonNode> = JacksonRuntime(configuration)
// Now the function is available in expressions
    val jackson = Jackson.defaultMapper()
    val inputTxt = """
    
    {
        "items": [
            { "id":"1", "foo":"barA", "type":"Type1"},
            { "id":"2", "foo":"barB", "type":"Type2"},
            { "id":"3", "foo":"barC", "type":"Type3"},
            { "id":"21", "foo":"barB", "type":"Type2"},
            { "id":"31", "foo":"barC", "type":"Type3"}
        ]
    }
    
    """
    val input = jackson.readTree(inputTxt)

    val q = runtime.compile("group_by(&type, items)")

    val result = q.search(input)

    println(result)
}


class GroupByFunction :
    BaseFunction(
        "group_by",
        ArgumentConstraints.expression(),
        ArgumentConstraints.arrayOf(ArgumentConstraints.anyValue())
    ) {
    // T:JsonNode
    override fun <T> callFunction(runtime: Adapter<T>, arguments: List<FunctionArgument<T>>): T {
        val keySelectorExpression = arguments[0].expression()
        val array = arguments[1].value()
        val arrayType = runtime.typeOf(array)
        if (arrayType != JmesPathType.ARRAY) {
            return runtime.createNull()
        }


        val elements = runtime.toList(array).toList()

        val grouped:Map<T,List<T>> = elements.groupBy { element ->
            val keySelector = keySelectorExpression.search(element)
            keySelector
        }

        val groupedNode = createObject(grouped)
        return groupedNode as T
    }

    private fun <T> createObject(obj: Map<T, List<T>>): JsonNode {
        val objectNode = ObjectNode(JsonNodeFactory.instance)
        for ((key, value) in obj) {
            key as JsonNode
            value as List<JsonNode>
            val keyTxt: String = key.textValue() ?: "null"
            val valueNode: JsonNode = createArray(value)
            objectNode.set<JsonNode>(keyTxt, valueNode)
        }
        return objectNode
    }

    private fun createArray(elements: Collection<JsonNode>): JsonNode {
        val array:ArrayNode = JsonNodeFactory.instance.arrayNode()
        array.addAll(elements)
        return array
    }
}