lance_index/vector/hnsw/

builder.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

//! Builder of Hnsw Graph.

use arrow::array::{AsArray, ListBuilder, UInt32Builder};
use arrow::compute::concat_batches;
use arrow::datatypes::{Float32Type, UInt32Type};
use arrow_array::{ArrayRef, Float32Array, RecordBatch, UInt64Array};
use crossbeam_queue::ArrayQueue;
use deepsize::DeepSizeOf;
use itertools::Itertools;

use lance_core::utils::tokio::get_num_compute_intensive_cpus;
use lance_linalg::distance::DistanceType;
use rayon::prelude::*;
use std::cmp::min;
use std::collections::{BinaryHeap, HashMap, VecDeque};
use std::fmt::Debug;
use std::iter;
use std::sync::Arc;
use std::sync::RwLock;
use std::sync::atomic::{AtomicUsize, Ordering};
use tracing::instrument;

use lance_core::{Error, Result};
use rand::{Rng, rng};
use serde::{Deserialize, Serialize};

use super::super::graph::beam_search;
use super::{HNSW_TYPE, HnswMetadata, VECTOR_ID_COL, VECTOR_ID_FIELD, select_neighbors_heuristic};
use crate::metrics::MetricsCollector;
use crate::prefilter::PreFilter;
use crate::vector::flat::storage::FlatFloatStorage;
use crate::vector::graph::builder::GraphBuilderNode;
use crate::vector::graph::{
    BorrowingGraph, DISTS_FIELD, Graph, NEIGHBORS_COL, NEIGHBORS_FIELD, OrderedFloat, OrderedNode,
    VisitedGenerator,
};
use crate::vector::graph::{Visited, beam_search_borrowed, greedy_search, greedy_search_borrowed};
use crate::vector::storage::{DistCalculator, VectorStore};
use crate::vector::v3::subindex::IvfSubIndex;
use crate::vector::{DIST_COL, Query, VECTOR_RESULT_SCHEMA};

pub const HNSW_METADATA_KEY: &str = "lance:hnsw";

/// Parameters of building HNSW index
#[derive(Debug, Clone, Serialize, Deserialize, DeepSizeOf)]
pub struct HnswBuildParams {
    /// max level ofm
    pub max_level: u16,

    /// number of connections to establish while inserting new element
    pub m: usize,

    /// size of the dynamic list for the candidates
    pub ef_construction: usize,

    /// number of vectors ahead to prefetch while building the graph
    pub prefetch_distance: Option<usize>,
}

impl Default for HnswBuildParams {
    fn default() -> Self {
        Self {
            max_level: 7,
            m: 20,
            ef_construction: 150,
            prefetch_distance: Some(2),
        }
    }
}

impl HnswBuildParams {
    /// The maximum level of the graph.
    /// The default value is `8`.
    pub fn max_level(mut self, max_level: u16) -> Self {
        self.max_level = max_level;
        self
    }

    /// The number of connections to establish while inserting new element
    /// The default value is `30`.
    pub fn num_edges(mut self, m: usize) -> Self {
        self.m = m;
        self
    }

    /// Number of candidates to be considered when searching for the nearest neighbors
    /// during the construction of the graph.
    ///
    /// The default value is `100`.
    pub fn ef_construction(mut self, ef_construction: usize) -> Self {
        self.ef_construction = ef_construction;
        self
    }

    /// Build the HNSW index from the given data.
    ///
    /// # Parameters
    /// - `data`: A FixedSizeList to build the HNSW.
    /// - `distance_type`: The distance type to use.
    pub async fn build(self, data: ArrayRef, distance_type: DistanceType) -> Result<HNSW> {
        let vec_store = Arc::new(FlatFloatStorage::new(
            data.as_fixed_size_list().clone(),
            distance_type,
        ));
        HNSW::index_vectors(vec_store.as_ref(), self)
    }
}

/// Build a HNSW graph.
///
/// Currently, the HNSW graph is fully built in memory.
///
/// During the build, the graph is built layer by layer.
///
/// Each node in the graph has a global ID which is the index on the base layer.
#[derive(Clone, DeepSizeOf)]
pub struct HNSW {
    inner: Arc<HnswCore>,
}

struct HnswCore {
    params: HnswBuildParams,
    nodes: Arc<Vec<GraphBuilderNode>>,
    level_count: Vec<usize>,
    entry_point: u32,
    visited_generator_queue: Arc<ArrayQueue<VisitedGenerator>>,
}

impl DeepSizeOf for HnswCore {
    fn deep_size_of_children(&self, context: &mut deepsize::Context) -> usize {
        self.params.deep_size_of_children(context)
            + self.nodes.deep_size_of_children(context)
            + self.level_count.deep_size_of_children(context)
        // Skipping the visited_generator_queue
    }
}

impl HnswCore {
    fn max_level(&self) -> u16 {
        self.params.max_level
    }

    fn num_nodes(&self, level: usize) -> usize {
        self.level_count[level]
    }

    fn nodes(&self) -> Arc<Vec<GraphBuilderNode>> {
        self.nodes.clone()
    }
}

impl Debug for HNSW {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "HNSW(max_layers: {})", self.inner.max_level() as usize,)
    }
}

impl HNSW {
    pub fn empty() -> Self {
        Self {
            inner: Arc::new(HnswCore {
                params: HnswBuildParams::default(),
                nodes: Arc::new(Vec::new()),
                level_count: Vec::new(),
                entry_point: 0,
                visited_generator_queue: Arc::new(ArrayQueue::new(1)),
            }),
        }
    }

    pub fn len(&self) -> usize {
        self.inner.nodes.len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn max_level(&self) -> u16 {
        self.inner.max_level()
    }

    pub fn num_nodes(&self, level: usize) -> usize {
        self.inner.num_nodes(level)
    }

    pub fn nodes(&self) -> Arc<Vec<GraphBuilderNode>> {
        self.inner.nodes()
    }

    #[allow(clippy::too_many_arguments)]
    pub fn search_inner(
        &self,
        query: ArrayRef,
        k: usize,
        params: &HnswQueryParams,
        bitset: Option<Visited>,
        visited_generator: &mut VisitedGenerator,
        storage: &impl VectorStore,
        prefetch_distance: Option<usize>,
    ) -> Result<Vec<OrderedNode>> {
        let dist_calc = storage.dist_calculator(query, params.dist_q_c);
        let mut ep = OrderedNode::new(0, dist_calc.distance(0).into());
        let nodes = self.inner.nodes.as_ref();
        for level in (0..self.max_level()).rev() {
            let cur_level = ImmutableHnswLevelView::new(level, nodes);
            ep = greedy_search_borrowed(
                &cur_level,
                ep,
                &dist_calc,
                self.inner.params.prefetch_distance,
            );
        }

        let bottom_level = ImmutableHnswBottomView::new(nodes);
        let mut visited = visited_generator.generate(storage.len());
        Ok(beam_search_borrowed(
            &bottom_level,
            &ep,
            params,
            &dist_calc,
            bitset.as_ref(),
            prefetch_distance,
            &mut visited,
        )
        .into_iter()
        .take(k)
        .collect())
    }

    #[instrument(level = "debug", skip(self, query, bitset, storage))]
    pub fn search_basic(
        &self,
        query: ArrayRef,
        k: usize,
        params: &HnswQueryParams,
        bitset: Option<Visited>,
        storage: &impl VectorStore,
    ) -> Result<Vec<OrderedNode>> {
        let mut visited_generator = self
            .inner
            .visited_generator_queue
            .pop()
            .unwrap_or_else(|| VisitedGenerator::new(storage.len()));
        let result = self.search_inner(
            query,
            k,
            params,
            bitset,
            &mut visited_generator,
            storage,
            Some(2),
        );

        match self.inner.visited_generator_queue.push(visited_generator) {
            Ok(_) => {}
            Err(_) => {
                log::warn!("visited_generator_queue is full");
            }
        }

        result
    }

    #[instrument(level = "debug", skip(self, storage, query, prefilter_bitset))]
    fn flat_search(
        &self,
        storage: &impl VectorStore,
        query: ArrayRef,
        k: usize,
        prefilter_bitset: Visited,
        params: &HnswQueryParams,
    ) -> Vec<OrderedNode> {
        let lower_bound: OrderedFloat = params.lower_bound.unwrap_or(f32::MIN).into();
        let upper_bound: OrderedFloat = params.upper_bound.unwrap_or(f32::MAX).into();

        let dist_calc = storage.dist_calculator(query, params.dist_q_c);
        let mut heap = BinaryHeap::<OrderedNode>::with_capacity(k);

        match self.inner.params.prefetch_distance {
            Some(ahead) if ahead > 0 => {
                let mut ids_iter = prefilter_bitset.iter_ones().map(|i| i as u32);
                let mut buffer = VecDeque::with_capacity(ahead + 1);
                for _ in 0..=ahead {
                    if let Some(id) = ids_iter.next() {
                        buffer.push_back(id);
                    } else {
                        break;
                    }
                }

                while let Some(node_id) = buffer.pop_front() {
                    if let Some(&prefetch_id) = buffer.get(ahead - 1) {
                        dist_calc.prefetch(prefetch_id);
                    }
                    if let Some(next) = ids_iter.next() {
                        buffer.push_back(next);
                    }

                    let dist: OrderedFloat = dist_calc.distance(node_id).into();
                    if dist <= lower_bound || dist > upper_bound {
                        continue;
                    }
                    if heap.len() < k {
                        heap.push((dist, node_id).into());
                    } else if dist < heap.peek().unwrap().dist {
                        heap.pop();
                        heap.push((dist, node_id).into());
                    }
                }
            }
            _ => {
                for node_id in prefilter_bitset.iter_ones().map(|i| i as u32) {
                    let dist: OrderedFloat = dist_calc.distance(node_id).into();
                    if dist <= lower_bound || dist > upper_bound {
                        continue;
                    }
                    if heap.len() < k {
                        heap.push((dist, node_id).into());
                    } else if dist < heap.peek().unwrap().dist {
                        heap.pop();
                        heap.push((dist, node_id).into());
                    }
                }
            }
        };
        heap.into_sorted_vec()
    }

    /// Returns the metadata of this [`HNSW`].
    pub fn metadata(&self) -> HnswMetadata {
        // calculate the offsets of each level,
        // start from 0
        let level_offsets = self
            .inner
            .level_count
            .iter()
            .chain(iter::once(&0))
            .scan(0, |state, x| {
                let start = *state;
                *state += *x;
                Some(start)
            })
            .collect();

        HnswMetadata {
            entry_point: self.inner.entry_point,
            params: self.inner.params.clone(),
            level_offsets,
        }
    }
}

struct HnswBuilder {
    params: HnswBuildParams,

    nodes: Arc<Vec<RwLock<GraphBuilderNode>>>,
    level_count: Vec<AtomicUsize>,

    entry_point: u32,

    visited_generator_queue: Arc<ArrayQueue<VisitedGenerator>>,
}

impl DeepSizeOf for HnswBuilder {
    fn deep_size_of_children(&self, context: &mut deepsize::Context) -> usize {
        self.params.deep_size_of_children(context)
            + self.nodes.deep_size_of_children(context)
            + self.level_count.deep_size_of_children(context)
        // Skipping the visited_generator_queue
    }
}

impl HnswBuilder {
    fn finish(self) -> HNSW {
        let nodes = match Arc::try_unwrap(self.nodes) {
            Ok(nodes) => nodes
                .into_iter()
                .map(|node| node.into_inner().expect("builder lock poisoned"))
                .collect(),
            Err(nodes) => nodes
                .iter()
                .map(|node| node.read().expect("builder lock poisoned").clone())
                .collect(),
        };

        let level_count = self
            .level_count
            .into_iter()
            .map(|count| count.load(Ordering::Relaxed))
            .collect();

        HNSW {
            inner: Arc::new(HnswCore {
                params: self.params,
                nodes: Arc::new(nodes),
                level_count,
                entry_point: self.entry_point,
                visited_generator_queue: self.visited_generator_queue,
            }),
        }
    }

    /// Create a new [`HNSWBuilder`] with prepared params and in memory vector storage.
    pub fn with_params(params: HnswBuildParams, storage: &impl VectorStore) -> Self {
        let len = storage.len();
        let max_level = params.max_level;

        let level_count = (0..max_level)
            .map(|_| AtomicUsize::new(0))
            .collect::<Vec<_>>();

        let visited_generator_queue = Arc::new(ArrayQueue::new(get_num_compute_intensive_cpus()));
        for _ in 0..get_num_compute_intensive_cpus() {
            visited_generator_queue
                .push(VisitedGenerator::new(0))
                .unwrap();
        }
        let mut builder = Self {
            params,
            nodes: Arc::new(Vec::new()),
            level_count,
            entry_point: 0,
            visited_generator_queue,
        };

        if storage.is_empty() {
            return builder;
        }

        let mut nodes = Vec::with_capacity(len);
        {
            if len > 0 {
                nodes.push(RwLock::new(GraphBuilderNode::new(0, max_level as usize)));
            }
            let mut level_rng = rng();
            for i in 1..len {
                nodes.push(RwLock::new(GraphBuilderNode::new(
                    i as u32,
                    builder.random_level(&mut level_rng) as usize + 1,
                )));
            }
        }
        builder.nodes = Arc::new(nodes);

        builder
    }

    /// New node's level
    ///
    /// See paper `Algorithm 1`
    fn random_level<R: Rng + ?Sized>(&self, rng: &mut R) -> u16 {
        let ml = 1.0 / (self.params.m as f32).ln();
        min(
            (-rng.random::<f32>().ln() * ml) as u16,
            self.params.max_level - 1,
        )
    }

    /// Insert one node.
    fn insert(
        &self,
        node: u32,
        visited_generator: &mut VisitedGenerator,
        storage: &impl VectorStore,
    ) {
        let nodes = &self.nodes;
        let target_level = nodes[node as usize].read().unwrap().level_neighbors.len() as u16 - 1;
        let dist_calc = storage.dist_calculator_from_id(node);
        let mut ep = OrderedNode::new(
            self.entry_point,
            dist_calc.distance(self.entry_point).into(),
        );

        //
        // Search for entry point in paper.
        // ```
        //   for l_c in (L..l+1) {
        //     W = Search-Layer(q, ep, ef=1, l_c)
        //    ep = Select-Neighbors(W, 1)
        //  }
        // ```
        for level in (target_level + 1..self.params.max_level).rev() {
            let cur_level = HnswLevelView::new(level, nodes);
            ep = greedy_search(&cur_level, ep, &dist_calc, self.params.prefetch_distance);
        }

        let mut pruned_neighbors_per_level: Vec<Vec<_>> =
            vec![Vec::new(); (target_level + 1) as usize];
        {
            let mut current_node = nodes[node as usize].write().unwrap();
            for level in (0..=target_level).rev() {
                self.level_count[level as usize].fetch_add(1, Ordering::Relaxed);

                let neighbors = self.search_level(&ep, level, &dist_calc, nodes, visited_generator);
                for neighbor in &neighbors {
                    current_node.add_neighbor(neighbor.id, neighbor.dist, level);
                }
                self.prune(storage, &mut current_node, level);
                pruned_neighbors_per_level[level as usize]
                    .clone_from(&current_node.level_neighbors_ranked[level as usize]);

                ep = neighbors[0].clone();
            }
        }
        for (level, pruned_neighbors) in pruned_neighbors_per_level.iter().enumerate() {
            let _: Vec<_> = pruned_neighbors
                .iter()
                .map(|unpruned_edge| {
                    let level = level as u16;
                    let m_max = match level {
                        0 => self.params.m * 2,
                        _ => self.params.m,
                    };
                    if unpruned_edge.dist
                        < nodes[unpruned_edge.id as usize]
                            .read()
                            .unwrap()
                            .cutoff(level, m_max)
                    {
                        let mut chosen_node = nodes[unpruned_edge.id as usize].write().unwrap();
                        chosen_node.add_neighbor(node, unpruned_edge.dist, level);
                        self.prune(storage, &mut chosen_node, level);
                    }
                })
                .collect();
        }
    }

    fn search_level(
        &self,
        ep: &OrderedNode,
        level: u16,
        dist_calc: &impl DistCalculator,
        nodes: &[RwLock<GraphBuilderNode>],
        visited_generator: &mut VisitedGenerator,
    ) -> Vec<OrderedNode> {
        let cur_level = HnswLevelView::new(level, nodes);
        let mut visited = visited_generator.generate(nodes.len());
        beam_search(
            &cur_level,
            ep,
            &HnswQueryParams {
                ef: self.params.ef_construction,
                lower_bound: None,
                upper_bound: None,
                dist_q_c: 0.0,
            },
            dist_calc,
            None,
            self.params.prefetch_distance,
            &mut visited,
        )
    }

    fn prune(&self, storage: &impl VectorStore, builder_node: &mut GraphBuilderNode, level: u16) {
        let m_max = match level {
            0 => self.params.m * 2,
            _ => self.params.m,
        };

        let neighbors_ranked = &mut builder_node.level_neighbors_ranked[level as usize];
        let level_neighbors = neighbors_ranked.clone();
        if level_neighbors.len() <= m_max {
            builder_node.update_from_ranked_neighbors(level);
            return;
        }

        *neighbors_ranked = select_neighbors_heuristic(storage, &level_neighbors, m_max);
        builder_node.update_from_ranked_neighbors(level);
    }
}

// View of a level in HNSW graph.
// This is used to iterate over neighbors in a specific level.
pub(crate) struct HnswLevelView<'a> {
    level: u16,
    nodes: &'a [RwLock<GraphBuilderNode>],
}

impl<'a> HnswLevelView<'a> {
    pub fn new(level: u16, nodes: &'a [RwLock<GraphBuilderNode>]) -> Self {
        Self { level, nodes }
    }
}

impl Graph for HnswLevelView<'_> {
    fn len(&self) -> usize {
        self.nodes.len()
    }

    fn neighbors(&self, key: u32) -> Arc<Vec<u32>> {
        let node = &self.nodes[key as usize];
        node.read().unwrap().level_neighbors[self.level as usize].clone()
    }
}

pub(crate) struct ImmutableHnswLevelView<'a> {
    level: u16,
    nodes: &'a [GraphBuilderNode],
}

impl<'a> ImmutableHnswLevelView<'a> {
    pub fn new(level: u16, nodes: &'a [GraphBuilderNode]) -> Self {
        Self { level, nodes }
    }
}

impl Graph for ImmutableHnswLevelView<'_> {
    fn len(&self) -> usize {
        self.nodes.len()
    }

    fn neighbors(&self, key: u32) -> Arc<Vec<u32>> {
        self.nodes[key as usize].level_neighbors[self.level as usize].clone()
    }
}

impl BorrowingGraph for ImmutableHnswLevelView<'_> {
    fn len(&self) -> usize {
        self.nodes.len()
    }

    fn neighbors(&self, key: u32) -> &[u32] {
        self.nodes[key as usize].level_neighbors[self.level as usize].as_slice()
    }
}

pub(crate) struct ImmutableHnswBottomView<'a> {
    nodes: &'a [GraphBuilderNode],
}

impl<'a> ImmutableHnswBottomView<'a> {
    pub fn new(nodes: &'a [GraphBuilderNode]) -> Self {
        Self { nodes }
    }
}

impl Graph for ImmutableHnswBottomView<'_> {
    fn len(&self) -> usize {
        self.nodes.len()
    }

    fn neighbors(&self, key: u32) -> Arc<Vec<u32>> {
        self.nodes[key as usize].bottom_neighbors.clone()
    }
}

impl BorrowingGraph for ImmutableHnswBottomView<'_> {
    fn len(&self) -> usize {
        self.nodes.len()
    }

    fn neighbors(&self, key: u32) -> &[u32] {
        self.nodes[key as usize].bottom_neighbors.as_slice()
    }
}

#[derive(Debug, Clone, Copy)]
pub struct HnswQueryParams {
    pub ef: usize,
    pub lower_bound: Option<f32>,
    pub upper_bound: Option<f32>,
    pub dist_q_c: f32,
}

impl From<&Query> for HnswQueryParams {
    fn from(query: &Query) -> Self {
        let k = query.k * query.refine_factor.unwrap_or(1) as usize;
        Self {
            ef: query.ef.unwrap_or(k + k / 2),
            lower_bound: query.lower_bound,
            upper_bound: query.upper_bound,
            dist_q_c: query.dist_q_c,
        }
    }
}

impl IvfSubIndex for HNSW {
    type BuildParams = HnswBuildParams;
    type QueryParams = HnswQueryParams;

    fn load(data: RecordBatch) -> Result<Self>
    where
        Self: Sized,
    {
        if data.num_rows() == 0 {
            return Ok(Self::empty());
        }

        let hnsw_metadata = data
            .schema_ref()
            .metadata()
            .get(HNSW_METADATA_KEY)
            .ok_or(Error::index(format!("{} not found", HNSW_METADATA_KEY)))?;
        let hnsw_metadata: HnswMetadata = serde_json::from_str(hnsw_metadata).map_err(|e| {
            Error::index(format!(
                "Failed to decode HNSW metadata: {}, json: {}",
                e, hnsw_metadata
            ))
        })?;

        let levels: Vec<_> = hnsw_metadata
            .level_offsets
            .iter()
            .tuple_windows()
            .map(|(start, end)| data.slice(*start, end - start))
            .collect();

        let level_count = levels.iter().map(|b| b.num_rows()).collect::<Vec<_>>();

        let bottom_level_len = levels[0].num_rows();
        let mut nodes = Vec::with_capacity(bottom_level_len);
        for i in 0..bottom_level_len {
            nodes.push(GraphBuilderNode::new(i as u32, levels.len()));
        }
        for (level, batch) in levels.into_iter().enumerate() {
            let ids = batch[VECTOR_ID_COL].as_primitive::<UInt32Type>();
            let neighbors = batch[NEIGHBORS_COL].as_list::<i32>();
            let distances = batch[DIST_COL].as_list::<i32>();

            for ((node, neighbors), distances) in
                ids.iter().zip(neighbors.iter()).zip(distances.iter())
            {
                let node = node.unwrap();
                let neighbors = neighbors.as_ref().unwrap().as_primitive::<UInt32Type>();
                let distances = distances.as_ref().unwrap().as_primitive::<Float32Type>();

                nodes[node as usize].level_neighbors_ranked[level] = neighbors
                    .iter()
                    .zip(distances.iter())
                    .map(|(n, dist)| OrderedNode::new(n.unwrap(), OrderedFloat(dist.unwrap())))
                    .collect();
                nodes[node as usize].update_from_ranked_neighbors(level as u16);
            }
        }

        let visited_generator_queue =
            Arc::new(ArrayQueue::new(get_num_compute_intensive_cpus() * 2));
        for _ in 0..get_num_compute_intensive_cpus() * 2 {
            visited_generator_queue
                .push(VisitedGenerator::new(0))
                .unwrap();
        }
        let inner = HnswCore {
            params: hnsw_metadata.params,
            nodes: Arc::new(nodes),
            level_count,
            entry_point: hnsw_metadata.entry_point,
            visited_generator_queue,
        };

        Ok(Self {
            inner: Arc::new(inner),
        })
    }

    fn name() -> &'static str {
        HNSW_TYPE
    }

    fn metadata_key() -> &'static str {
        "lance:hnsw"
    }

    /// Return the schema of the sub index
    fn schema() -> arrow_schema::SchemaRef {
        arrow_schema::Schema::new(vec![
            VECTOR_ID_FIELD.clone(),
            NEIGHBORS_FIELD.clone(),
            DISTS_FIELD.clone(),
        ])
        .into()
    }

    #[instrument(level = "debug", skip(self, query, storage, prefilter, _metrics))]
    fn search(
        &self,
        query: ArrayRef,
        k: usize,
        params: Self::QueryParams,
        storage: &impl VectorStore,
        prefilter: Arc<dyn PreFilter>,
        _metrics: &dyn MetricsCollector,
    ) -> Result<RecordBatch> {
        if params.ef < k {
            return Err(Error::index(
                "ef must be greater than or equal to k".to_string(),
            ));
        }

        let schema = VECTOR_RESULT_SCHEMA.clone();
        if self.is_empty() {
            return Ok(RecordBatch::new_empty(schema));
        }

        let mut prefilter_generator = self
            .inner
            .visited_generator_queue
            .pop()
            .unwrap_or_else(|| VisitedGenerator::new(storage.len()));
        let prefilter_bitset = if prefilter.is_empty() {
            None
        } else {
            let indices = prefilter.filter_row_ids(Box::new(storage.row_ids()));
            let mut bitset = prefilter_generator.generate(storage.len());
            for indices in indices {
                bitset.insert(indices as u32);
            }
            Some(bitset)
        };

        let remained = prefilter_bitset
            .as_ref()
            .map(|b| b.count_ones())
            .unwrap_or(storage.len());
        let results = if remained < self.len() * 10 / 100 {
            let prefilter_bitset =
                prefilter_bitset.expect("the prefilter bitset must be set for flat search");
            self.flat_search(storage, query, k, prefilter_bitset, &params)
        } else {
            self.search_basic(query, k, &params, prefilter_bitset, storage)?
        };
        // if the queue is full, we just don't push it back, so ignore the error here
        let _ = self.inner.visited_generator_queue.push(prefilter_generator);

        // need to unique by row ids in case of searching multivector
        let (row_ids, dists): (Vec<_>, Vec<_>) = results
            .into_iter()
            .map(|r| (storage.row_id(r.id), r.dist.0))
            .unique_by(|r| r.0)
            .unzip();
        let row_ids = Arc::new(UInt64Array::from(row_ids));
        let distances = Arc::new(Float32Array::from(dists));

        Ok(RecordBatch::try_new(schema, vec![distances, row_ids])?)
    }

    /// Given a vector storage, containing all the data for the IVF partition, build the sub index.
    fn index_vectors(storage: &impl VectorStore, params: Self::BuildParams) -> Result<Self>
    where
        Self: Sized,
    {
        let builder = HnswBuilder::with_params(params, storage);

        log::debug!(
            "Building HNSW graph: num={}, max_levels={}, m={}, ef_construction={}, distance_type:{}",
            storage.len(),
            builder.params.max_level,
            builder.params.m,
            builder.params.ef_construction,
            storage.distance_type(),
        );

        if storage.is_empty() {
            return Ok(builder.finish());
        }

        let len = storage.len();
        builder.level_count[0].fetch_add(1, Ordering::Relaxed);
        (1..len).into_par_iter().for_each_init(
            || VisitedGenerator::new(len),
            |visited_generator, node| {
                builder.insert(node as u32, visited_generator, storage);
            },
        );

        assert_eq!(builder.level_count[0].load(Ordering::Relaxed), len);
        Ok(builder.finish())
    }

    fn remap(
        &self,
        _mapping: &HashMap<u64, Option<u64>>, // we don't need the mapping here because we rebuild the graph from remapped storage
        store: &impl VectorStore,
    ) -> Result<Self> {
        // We can't simply remap the row ids in the graph because the vectors are changed,
        // so the graph needs to be rebuilt.
        Self::index_vectors(store, self.inner.params.clone())
    }

    /// Encode the sub index into a record batch
    fn to_batch(&self) -> Result<RecordBatch> {
        let mut vector_id_builder = UInt32Builder::with_capacity(self.len());
        let mut neighbors_builder = ListBuilder::with_capacity(UInt32Builder::new(), self.len());
        let mut distances_builder =
            ListBuilder::with_capacity(arrow_array::builder::Float32Builder::new(), self.len());
        let mut batches = Vec::with_capacity(self.max_level() as usize);
        for level in 0..self.max_level() {
            let level = level as usize;
            for (id, node) in self.inner.nodes.iter().enumerate() {
                if level >= node.level_neighbors.len() {
                    continue;
                }
                let neighbors = node.level_neighbors[level].iter().map(|n| Some(*n));
                let distances = node.level_neighbors_ranked[level]
                    .iter()
                    .map(|n| Some(n.dist.0));
                vector_id_builder.append_value(id as u32);
                neighbors_builder.append_value(neighbors);
                distances_builder.append_value(distances);
            }

            let batch = RecordBatch::try_new(
                Self::schema(),
                vec![
                    Arc::new(vector_id_builder.finish()),
                    Arc::new(neighbors_builder.finish()),
                    Arc::new(distances_builder.finish()),
                ],
            )?;
            batches.push(batch);
        }

        let metadata = self.metadata();
        let metadata = serde_json::to_string(&metadata)?;
        let schema = Self::schema()
            .as_ref()
            .clone()
            .with_metadata(HashMap::from_iter(vec![(
                HNSW_METADATA_KEY.to_string(),
                metadata,
            )]));
        let batch = concat_batches(&Self::schema(), batches.iter())?;
        let batch = batch.with_schema(Arc::new(schema))?;
        Ok(batch)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use arrow_array::FixedSizeListArray;
    use arrow_schema::Schema;
    use lance_arrow::FixedSizeListArrayExt;
    use lance_file::previous::{
        reader::FileReader as PreviousFileReader,
        writer::{
            FileWriter as PreviousFileWriter, FileWriterOptions as PreviousFileWriterOptions,
        },
    };
    use lance_io::object_store::ObjectStore;
    use lance_linalg::distance::DistanceType;
    use lance_table::format::SelfDescribingFileReader;
    use lance_table::io::manifest::ManifestDescribing;
    use lance_testing::datagen::generate_random_array;
    use object_store::path::Path;

    use crate::scalar::IndexWriter;
    use crate::vector::v3::subindex::IvfSubIndex;
    use crate::vector::{
        flat::storage::FlatFloatStorage,
        graph::{DISTS_FIELD, NEIGHBORS_FIELD},
        hnsw::{
            HNSW, VECTOR_ID_FIELD,
            builder::{HnswBuildParams, HnswQueryParams},
        },
    };

    #[tokio::test]
    async fn test_builder_write_load() {
        const DIM: usize = 32;
        const TOTAL: usize = 2048;
        const NUM_EDGES: usize = 20;
        let data = generate_random_array(TOTAL * DIM);
        let fsl = FixedSizeListArray::try_new_from_values(data, DIM as i32).unwrap();
        let store = Arc::new(FlatFloatStorage::new(fsl.clone(), DistanceType::L2));
        let builder = HNSW::index_vectors(
            store.as_ref(),
            HnswBuildParams::default()
                .num_edges(NUM_EDGES)
                .ef_construction(50),
        )
        .unwrap();

        let object_store = ObjectStore::memory();
        let path = Path::from("test_builder_write_load");
        let writer = object_store.create(&path).await.unwrap();
        let schema = Schema::new(vec![
            VECTOR_ID_FIELD.clone(),
            NEIGHBORS_FIELD.clone(),
            DISTS_FIELD.clone(),
        ]);
        let schema = lance_core::datatypes::Schema::try_from(&schema).unwrap();
        let mut writer = PreviousFileWriter::<ManifestDescribing>::with_object_writer(
            writer,
            schema,
            &PreviousFileWriterOptions::default(),
        )
        .unwrap();
        let batch = builder.to_batch().unwrap();
        let metadata = batch.schema_ref().metadata().clone();
        writer.write_record_batch(batch).await.unwrap();
        writer.finish_with_metadata(&metadata).await.unwrap();

        let reader = PreviousFileReader::try_new_self_described(&object_store, &path, None)
            .await
            .unwrap();
        let batch = reader
            .read_range(0..reader.len(), reader.schema())
            .await
            .unwrap();
        let loaded_hnsw = HNSW::load(batch).unwrap();

        let query = fsl.value(0);
        let k = 10;
        let params = HnswQueryParams {
            ef: 50,
            lower_bound: None,
            upper_bound: None,
            dist_q_c: 0.0,
        };
        let builder_results = builder
            .search_basic(query.clone(), k, &params, None, store.as_ref())
            .unwrap();
        let loaded_results = loaded_hnsw
            .search_basic(query, k, &params, None, store.as_ref())
            .unwrap();
        assert_eq!(builder_results, loaded_results);
    }
}