Получить первые N элементов из массива в таблице BigQuery

Question

У меня есть столбец массива, и я хотел бы получить первые N его элементы (сохраняя тип данных массива). Есть ли какой-нибудь хороший способ, как это сделать? В идеале без удаления, ранжирования и array_agg обратно в массив.

Я также мог бы сделать это (для получения первых 2 элементов):

WITH data AS
(
  SELECT 1001 as id, ['a', 'b', 'c'] as array_1
  UNION ALL
  SELECT 1002 as id, ['d', 'e', 'f', 'g'] as array_1
  UNION ALL
  SELECT 1003 as id, ['h', 'i'] as array_1
)
select *,
       [array_1[SAFE_OFFSET(0)], array_1[SAFE_OFFSET(1)]] as my_result
from data

Но, очевидно, это не очень хорошее решение, так какошибка в случае, если в каком-либо массиве будет только 1 элемент.

Answer 1

Вот общее решение с UDF, которое вы можете вызвать для любого типа массива:

CREATE TEMP FUNCTION TopN(arr ANY TYPE, n INT64) AS (
  ARRAY(SELECT x FROM UNNEST(arr) AS x WITH OFFSET off WHERE off < n ORDER BY off)
);

WITH data AS
(
  SELECT 1001 as id, ['a', 'b', 'c'] as array_1
  UNION ALL
  SELECT 1002 as id, ['d', 'e', 'f', 'g'] as array_1
  UNION ALL
  SELECT 1003 as id, ['h', 'i'] as array_1
)
select *, TopN(array_1, 2) AS my_result
from data

Используется unnest и функция массива, которая, как вам кажется, вы не хотели использовать, но онаимеет то преимущество, что он достаточно общий, так что вы можете передать ему любой массив.

Answer 2

Еще одна опция для BigQuery Standard SQL (с JS UDF)

#standardSQL
CREATE TEMP FUNCTION FirstN(arr ARRAY<STRING>, N FLOAT64)
RETURNS ARRAY<STRING> LANGUAGE js AS """ 
  return arr.slice(0, N);
""";
SELECT *, 
  FirstN(array_1, 3) AS my_result
FROM data