How can I return a single value in rego function?

I have the following rego:

retry_count_key := "retry"
get_retries_count(str) := {x |
    some i
    parts := split(str[i], ":")
    parts[0] == retry_count_key
    x := to_number(parts[1])
}
max_retry_label := get_retries_count(input.labels)

When I use an input like:

{
  "labels": [
    "retry:5",
    "test"
  ]
}

max_retry_label is always an array of numbers. Ideally, I'd like to just return a single (the first) number.

I get:

{
    "max_retry_label": [
        5
    ],
}

I want:

{
    "max_retry_label": 5
}

I can do something hacky like this, but I'd like to just have it self contained in the function mostly because this policy is huge and I dont want future devs to have to know they need to do this:

max_retry_label := sum(get_retries_count(input.labels))

You could move the sum call to your own function:

get_retries_count(str) := sum({x |
    some i
    parts := split(str[i], ":")
    parts[0] == retry_count_key
    x := to_number(parts[1])
})

Or if you only wanted the first value, return only that:

get_retries_count(str) := [x |
    some i
    parts := split(str[i], ":")
    parts[0] == retry_count_key
    x := to_number(parts[1])
][0]

(changed to be an array comprehension though, as unordered sets have no concept of a "first" item)

You could also get fancy with unification and some pattern matching :)

get_retries_count(str) := sum([to_number(x) |
    some i
    [retry_count_key, x] = split(str[i], ":")
])