BigQuery UDF getKey

query.sql

SELECT
  metadata.document_id,
  STRUCT(
    getKey(payload.histograms, "b") AS b
  ) as histograms
FROM UNNEST([
  STRUCT(
    STRUCT(
      "bcfa6d63-d0dc-4e4a-a086-8e3844fca9cc" AS document_id
    ) AS metadata,
    STRUCT(
      STRUCT(
        [
          STRUCT("a" AS key, "A" AS value),
          STRUCT("b" AS key, "B" AS value)
        ] AS key_value
      ) AS histograms
    ) AS payload
  )
])

result.json

[
  {
    "document_id": "bcfa6d63-d0dc-4e4a-a086-8e3844fca9cc", 
    "histograms": {
      "b": "B"
    }
  }
]

UDF.sql

CREATE TEMP FUNCTION getKey(map ANY TYPE, k ANY TYPE) AS (
  (
    SELECT key_value.value
    FROM UNNEST(map.key_value) AS key_value
    WHERE key_value.key = k
    LIMIT 1
  )
);

can be verified with cat UDF.sql query.sql | bq query --nouse_legacy_sql --format=prettyjson --quiet | diff -u - result.json