BigQuery - найти N ближайших векторов

Question

У меня есть таблица больших запросов, в которой есть столбец с повторяющимся типом данных с 512-мерным вектором (с плавающей точкой).

Я хотел бы выполнить запрос, который находит N наиболее похожих векторов.

В моем случае сходство можно просто определить как внутренний продукт целевого вектора и каждого вектора в базе данных.

Я нашел и выполнил приведенный ниже запрос, который генерирует его для всех комбинаций втаблица:

#standardSQL
CREATE TABLE ml.url_cosine_similarity AS
WITH pairwise AS (
  SELECT t1.url AS id_1, t2.url AS id_2
  FROM `project.dataset.table` t1
  INNER JOIN `project.dataset.table` t2
  ON t1.url < t2.url
)
SELECT id_1, id_2, ( 
  SELECT 
    SUM(value1 * value2)/ 
    SQRT(SUM(value1 * value1))/ 
    SQRT(SUM(value2 * value2))
  FROM UNNEST(a.page_vector) value1 WITH OFFSET pos1 
  JOIN UNNEST(b.page_vector) value2 WITH OFFSET pos2 
  ON pos1 = pos2
  ) cosine_similarity
FROM pairwise t
JOIN `project.dataset.table` a ON a.url = id_1
JOIN `project.dataset.table` b ON b.url = id_2

Однако, поскольку я не совсем понимаю, как работают массивы в больших запросах, я не уверен, как изменить этот запрос, чтобы получить целевой вектор и вернуть N соседей.

Answer 1

См. Упрощенный пример - он возвращает 3 верхние ближайшие пары векторов в таблице

#standardSQL
WITH `project.dataset.table` AS (
  SELECT 1 id, [1,2,3,4,5] page_vector UNION ALL
  SELECT 2, [1,3,4,5,16] UNION ALL
  SELECT 3, [2,3,4,5,6] UNION ALL
  SELECT 4, [2,4,6,8,9] UNION ALL
  SELECT 5, [1,3,4,5,16] UNION ALL
  SELECT 6, [11,12,13,14,15] 
)
SELECT a.id id1, b.id id2, ( 
  SELECT 
    SUM(value1 * value2)/ 
    SQRT(SUM(value1 * value1))/ 
    SQRT(SUM(value2 * value2))
  FROM UNNEST(a.page_vector) value1 WITH OFFSET pos1 
  JOIN UNNEST(b.page_vector) value2 WITH OFFSET pos2 
  ON pos1 = pos2
  ) cosine_similarity
FROM `project.dataset.table` a
JOIN `project.dataset.table` b
ON a.id < b.id
ORDER BY cosine_similarity DESC
LIMIT 3  

с выводом

Row id1 id2 cosine_similarity    
1   2   5   1.0  
2   1   4   0.9986422261219272   
3   3   4   0.9962894120648842    

Если вы хотите вывести ближайшие векторы (скажем, два) для каждого вектора в таблице - см. приведенный ниже пример

#standardSQL
WITH `project.dataset.table` AS (
  SELECT 1 id, [1,2,3,4,5] page_vector UNION ALL
  SELECT 2, [1,3,4,5,16] UNION ALL
  SELECT 3, [2,3,4,5,6] UNION ALL
  SELECT 4, [2,4,6,8,9] UNION ALL
  SELECT 5, [1,3,4,5,16] UNION ALL
  SELECT 6, [11,12,13,14,15] 
)
SELECT id, ANY_VALUE(page_vector) page_vector, 
  ARRAY_AGG(
    STRUCT(id2 AS id, page_vector2 AS page_vector, cosine_similarity AS cosine_similarity) 
    ORDER BY cosine_similarity DESC 
    LIMIT 2
  ) similar_vectors
FROM (
  SELECT a.id, a.page_vector, 
    b.id id2, b.page_vector page_vector2, ( 
    SELECT 
      SUM(value1 * value2)/ 
      SQRT(SUM(value1 * value1))/ 
      SQRT(SUM(value2 * value2))
    FROM UNNEST(a.page_vector) value1 WITH OFFSET pos1 
    JOIN UNNEST(b.page_vector) value2 WITH OFFSET pos2 
    ON pos1 = pos2
    ) cosine_similarity
  FROM `project.dataset.table` a
  JOIN `project.dataset.table` b
  ON a.id != b.id
)
GROUP BY id
ORDER BY id  

, при этом будет получено значение ниже