Skip to content

Instantly share code, notes, and snippets.

@bastman
Created July 14, 2021 10:41
Show Gist options
  • Save bastman/fa038295dda8a92b06cd64650dbf4df6 to your computer and use it in GitHub Desktop.
Save bastman/fa038295dda8a92b06cd64650dbf4df6 to your computer and use it in GitHub Desktop.
groupBy impl for JMESPath (proof-of-concept)
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
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment