qdrant-edge 0.8.0

use std::borrow::Cow;
use std::cmp;
use std::collections::{HashMap, HashSet};
use std::hash::{Hash, Hasher};
use std::ops::Deref;
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;

use ahash::AHasher;
use atomic_refcell::AtomicRefCell;
use bitvec::macros::internal::funty::Integral;
use crate::common::budget::ResourcePermit;
use crate::common::counter::hardware_counter::HardwareCounterCell;
use crate::common::flags::feature_flags;
use crate::common::progress_tracker::ProgressTracker;
use crate::common::small_uint::U24;
use crate::common::storage_version::StorageVersion;
use crate::common::types::{DeferredBehavior, PointOffsetType};
use crate::common::universal_io::MmapFs;
use fs_err as fs;
use itertools::Itertools;
use rand::Rng;
use tempfile::TempDir;
use uuid::Uuid;

use super::{
    create_mutable_id_tracker, create_payload_storage, create_sparse_vector_storage,
    get_payload_index_path, get_vector_index_path, get_vector_storage_path,
    open_or_create_sparse_vector_index, open_vector_storage,
};
use crate::segment::common::error_logging::LogError;
use crate::segment::common::operation_error::{OperationError, OperationResult, check_process_stopped};
use crate::segment::entry::entry_point::StorageSegmentEntry as _;
use crate::segment::id_tracker::compressed::compressed_point_mappings::CompressedPointMappings;
use crate::segment::id_tracker::disk_id_tracker::DiskIdTracker;
use crate::segment::id_tracker::immutable_id_tracker::ImmutableIdTracker;
use crate::segment::id_tracker::in_memory_id_tracker::InMemoryIdTracker;
use crate::segment::id_tracker::{IdTracker, IdTrackerEnum, IdTrackerRead, for_each_unique_point};
use crate::segment::index::field_index::FieldIndex;
use crate::segment::index::sparse_index::sparse_vector_index::SparseVectorIndexOpenArgs;
use crate::segment::index::struct_payload_index::{IndexLoadMode, StorageType, StructPayloadIndex};
use crate::segment::index::{PayloadIndex, PayloadIndexRead, VectorIndexEnum};
use crate::segment::payload_storage::PayloadStorage;
use crate::segment::payload_storage::payload_storage_enum::PayloadStorageEnum;
use crate::segment::segment::{Segment, SegmentVersion};
use crate::segment::segment_constructor::batched_reader::{PointData, merge_from};
use crate::segment::segment_constructor::{
    VectorIndexBuildArgs, VectorIndexOpenArgs, build_vector_index, load_segment,
};
use crate::segment::types::{
    CompactExtendedPointId, ExtendedPointId, HnswGlobalConfig, PayloadFieldSchema, PayloadKeyType,
    SegmentConfig, SegmentState, SeqNumberType, VectorNameBuf,
};
use crate::segment::vector_storage::quantized::quantized_vectors::{
    QuantizedVectors, QuantizedVectorsStorageType,
};
use crate::segment::vector_storage::{VectorStorage, VectorStorageEnum, VectorStorageRead};

/// Structure for constructing segment out of several other segments
pub struct SegmentBuilder {
    version: SeqNumberType,
    id_tracker: IdTrackerEnum,
    payload_storage: PayloadStorageEnum,
    vector_data: HashMap<VectorNameBuf, VectorData>,
    segment_config: SegmentConfig,
    hnsw_global_config: HnswGlobalConfig,

    // The temporary segment directory
    temp_dir: TempDir,
    indexed_fields: HashMap<PayloadKeyType, PayloadFieldSchema>,

    // Payload key to defragment data to
    defragment_keys: Vec<PayloadKeyType>,

    // Vector names that currently exist in the live collection schema. Used to decide what to do
    // with a source vector name that is absent from this builder's (frozen) target schema:
    // - present here (or `None`, i.e. unknown) => cancel the merge (the CreateVectorName race);
    // - absent here => a previously deleted vector whose data still lingers in older segment
    //   files, safe to prune. See the merge loop in `update`.
    live_vector_names: Option<HashSet<VectorNameBuf>>,
}

struct VectorData {
    vector_storage: VectorStorageEnum,
    old_indices: Vec<Arc<AtomicRefCell<VectorIndexEnum>>>,
}

impl SegmentBuilder {
    pub fn new(
        temp_dir: &Path,
        segment_config: &SegmentConfig,
        hnsw_global_config: &HnswGlobalConfig,
    ) -> OperationResult<Self> {
        let temp_dir = create_temp_dir(temp_dir)?;

        let id_tracker = if segment_config.is_appendable() {
            // Deferred state is applied when the freshly built segment is reloaded
            // via `load_segment`. The transient builder tracker doesn't need it.
            IdTrackerEnum::MutableIdTracker(create_mutable_id_tracker(temp_dir.path(), None)?)
        } else {
            IdTrackerEnum::InMemoryIdTracker(InMemoryIdTracker::new())
        };

        let payload_storage = create_payload_storage(temp_dir.path(), segment_config)?;

        let mut vector_data = HashMap::new();

        for (vector_name, vector_config) in &segment_config.vector_data {
            let vector_storage_path = get_vector_storage_path(temp_dir.path(), vector_name);
            let vector_storage = open_vector_storage(vector_config, &vector_storage_path)?;

            vector_data.insert(
                vector_name.to_owned(),
                VectorData {
                    vector_storage,
                    old_indices: Vec::new(),
                },
            );
        }

        for (vector_name, sparse_vector_config) in &segment_config.sparse_vector_data {
            let vector_storage_path = get_vector_storage_path(temp_dir.path(), vector_name);

            let vector_storage = create_sparse_vector_storage(
                &vector_storage_path,
                &sparse_vector_config.storage_type,
            )?;

            vector_data.insert(
                vector_name.to_owned(),
                VectorData {
                    vector_storage,
                    old_indices: Vec::new(),
                },
            );
        }

        Ok(SegmentBuilder {
            version: Default::default(), // default version is 0
            id_tracker,
            payload_storage,
            vector_data,
            segment_config: segment_config.clone(),
            hnsw_global_config: hnsw_global_config.clone(),
            temp_dir,
            indexed_fields: Default::default(),
            defragment_keys: vec![],
            live_vector_names: None,
        })
    }

    pub fn set_defragment_keys(&mut self, keys: Vec<PayloadKeyType>) {
        self.defragment_keys = keys;
    }

    /// Set the vector names that currently exist in the live collection schema.
    ///
    /// When set, [`SegmentBuilder::update`] prunes (rather than rejects) source vector names that
    /// are absent from both the target schema and this set, i.e. vectors that were deleted from the
    /// collection but whose data still lingers in older segment files. Source vector names that are
    /// still present in the live schema keep cancelling the merge, since those signal the
    /// CreateVectorName-vs-optimizer race rather than a genuine deletion. When unset, every
    /// source-superset mismatch cancels (the conservative default).
    pub fn set_live_vector_names(&mut self, names: HashSet<VectorNameBuf>) {
        self.live_vector_names = Some(names);
    }

    pub fn remove_indexed_field(&mut self, field: &PayloadKeyType) {
        self.indexed_fields.remove(field);
    }

    pub fn remove_index_field_if_incompatible(
        &mut self,
        field: &PayloadKeyType,
        schema: &PayloadFieldSchema,
    ) {
        if let Some(existing_schema) = self.indexed_fields.get(field)
            && existing_schema != schema
        {
            self.indexed_fields.remove(field);
        }
    }

    pub fn add_indexed_field(&mut self, field: PayloadKeyType, schema: PayloadFieldSchema) {
        self.indexed_fields.insert(field, schema);
    }

    /// Get ordering value from the payload index
    ///
    /// Ordering value is used to sort points to keep points with the same payload together
    /// Under the assumption that points are queried together, this will reduce the number of
    /// random disk reads.
    ///
    /// Note: This value doesn't guarantee strict ordering in ambiguous cases.
    ///       It should only be used in optimization purposes, not for correctness.
    fn _get_ordering_value(
        internal_id: PointOffsetType,
        indices: &[FieldIndex],
        hw_counter: &HardwareCounterCell,
    ) -> u64 {
        let mut ordering = 0;
        for payload_index in indices {
            match payload_index {
                FieldIndex::IntMapIndex(index) => {
                    if let Some(numbers) = index.get_values(internal_id, hw_counter) {
                        for number in numbers {
                            ordering = ordering.wrapping_add(*number as u64);
                        }
                    }
                    break;
                }
                FieldIndex::KeywordIndex(index) => {
                    if let Some(keywords) = index.get_values(internal_id, hw_counter) {
                        for keyword in keywords {
                            let mut hasher = AHasher::default();
                            keyword.hash(&mut hasher);
                            ordering = ordering.wrapping_add(hasher.finish());
                        }
                    }
                    break;
                }
                FieldIndex::IntIndex(index) => {
                    if let Some(numbers) = index.get_values(internal_id) {
                        for number in numbers {
                            ordering = ordering.wrapping_add(number as u64);
                        }
                    }
                    break;
                }
                FieldIndex::FloatIndex(index) => {
                    if let Some(numbers) = index.get_values(internal_id) {
                        for number in numbers {
                            // Bit-level conversion of f64 to u64 preserves ordering
                            // (for positive numbers)
                            //
                            // 0.001 -> 4562254508917369340
                            // 0.01  -> 4576918229304087675
                            // 0.05  -> 4587366580439587226
                            // 0.1   -> 4591870180066957722
                            // 1     -> 4607182418800017408
                            // 2     -> 4611686018427387904
                            // 10    -> 4621819117588971520
                            ordering = ordering.wrapping_add(number.to_bits());
                        }
                    }
                    break;
                }
                FieldIndex::DatetimeIndex(index) => {
                    if let Some(dates) = index.get_values(internal_id) {
                        for date in dates {
                            ordering = ordering.wrapping_add(date as u64);
                        }
                    }
                    break;
                }
                FieldIndex::UuidMapIndex(index) => {
                    if let Some(ids) = index.get_values(internal_id, hw_counter) {
                        uuid_hash(&mut ordering, ids.map(Cow::into_owned));
                    }
                    break;
                }
                FieldIndex::UuidIndex(index) => {
                    if let Some(ids) = index.get_values(internal_id) {
                        uuid_hash(&mut ordering, ids);
                    }
                    break;
                }
                FieldIndex::GeoIndex(_) => {}
                FieldIndex::FullTextIndex(_) => {}
                FieldIndex::BoolIndex(_) => {}
                FieldIndex::NullIndex(_) => {}
            }
        }
        ordering
    }

    /// Update current segment builder with all (not deleted) vectors and payload from `segments`.
    /// Also defragments if the `defragment_key` is set.
    /// However only points in the same call get defragmented and grouped together.
    /// Therefore this function should only be called once, unless this behavior is desired.
    ///
    /// # Result
    ///
    /// * `bool` - if `true` - data successfully added, if `false` - process was interrupted
    ///
    pub fn update(
        &mut self,
        segments: &[&Segment],
        stopped: &AtomicBool,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<bool> {
        if segments.is_empty() {
            return Ok(true);
        }

        if segments.len() > U24::MAX as usize {
            return Err(OperationError::service_error("Too many segments to update"));
        }

        let locked_id_trackers = segments.iter().map(|s| s.id_tracker.borrow()).collect_vec();
        let max_point_count = locked_id_trackers
            .iter()
            .map(|id_tracker| id_tracker.available_point_count())
            .max();
        let mut points_to_insert = Vec::with_capacity(max_point_count.unwrap_or_default());
        for_each_unique_point(locked_id_trackers.iter().map(|i| i.deref()), |item| {
            points_to_insert.push(PointData {
                external_id: CompactExtendedPointId::from(item.external_id),
                segment_index: U24::new_wrapped(item.tracker_index as u32),
                internal_id: item.internal_id,
                version: item.version,
                ordering: 0,
            });
        });
        drop(locked_id_trackers);

        let payloads: Vec<_> = segments.iter().map(|i| i.payload_index.borrow()).collect();

        for defragment_key in &self.defragment_keys {
            for point_data in &mut points_to_insert {
                let Some(payload_indices) = payloads[point_data.segment_index.get() as usize]
                    .field_indexes
                    .get(defragment_key)
                else {
                    continue;
                };

                point_data.ordering = point_data.ordering.wrapping_add(Self::_get_ordering_value(
                    point_data.internal_id,
                    payload_indices,
                    hw_counter,
                ));
            }
        }

        if !self.defragment_keys.is_empty() {
            points_to_insert.sort_unstable_by_key(|i| i.ordering);
        }

        let src_segment_max_version = segments.iter().map(|i| i.version()).max().unwrap();
        self.version = cmp::max(self.version, src_segment_max_version);

        let vector_storages: Vec<_> = segments.iter().map(|i| &i.vector_data).collect();

        // The merge loop below iterates `self.vector_data` (target), so any vector name present on
        // a source segment but absent from the target schema is silently dropped. Whether that drop
        // is correct depends on *why* the source carries a name the target lacks:
        //
        // - DeleteVectorName: the name was removed from the collection schema, but its data still
        //   lives in older segment files. Pruning it on rebuild is the desired recovery.
        // - CreateVectorName-vs-optimizer race: an optimizer launched with a pre-`CreateVectorName(V)`
        //   config observes sources that gained V mid-flight. Dropping V here would emit a merged
        //   segment without V at version >= V_opnum and break the next optimization round (whose
        //   refreshed config has V). This must be cancelled and retried.
        //
        // We tell the two apart with the live collection schema (`live_vector_names`), which is
        // authoritative because the schema is persisted before the op is applied to segments: a
        // deleted name is gone from the live schema, a freshly created one is present. When the live
        // schema is unknown (`None`) we conservatively cancel.
        //
        // Cancel via `Cancelled` rather than `ServiceError` so the optimization worker treats it as
        // a recoverable cancellation (logged at debug, tracker marked Cancelled, no shard-level
        // optimizer_errors, no RED status); the retry under the refreshed config merges cleanly.
        for vector_storage in &vector_storages {
            for source_vector_name in vector_storage.keys() {
                if self.vector_data.contains_key(source_vector_name) {
                    continue;
                }
                let deleted_from_schema = self
                    .live_vector_names
                    .as_ref()
                    .is_some_and(|live| !live.contains(source_vector_name));
                if deleted_from_schema {
                    log::debug!(
                        "Dropping vector name {source_vector_name} from source segment during \
                         optimization; it was deleted from the collection schema"
                    );
                } else {
                    return Err(OperationError::cancelled(format!(
                        "Cannot update from other segment because it has an extra \
                         vector name {source_vector_name} not in the target schema; \
                         retry after optimizer config refresh"
                    )));
                }
            }
        }

        let internal_range_start = self.id_tracker.available_point_count() as PointOffsetType;
        let internal_range_end = internal_range_start + points_to_insert.len() as PointOffsetType;

        let new_internal_range = internal_range_start..internal_range_end;

        for (vector_name, vector_data) in &mut self.vector_data {
            check_process_stopped(stopped)?;

            let other_vector_storages = vector_storages
                .iter()
                .map(|i| {
                    // Symmetric counterpart to the source-superset check above:
                    // when target has a vector name a source lacks, the
                    // optimizer-vs-`DeleteVectorName` race is the typical
                    // cause (V was removed from originals before the proxy
                    // wrap, but the frozen `target_config` still has V).
                    // Use `Cancelled` so the optimization worker treats this
                    // as a recoverable cancellation — no shard-level
                    // `optimizer_errors`, no RED status — and the next round
                    // with refreshed config merges cleanly.
                    let other_vector_data = i.get(vector_name).ok_or_else(|| {
                        OperationError::cancelled(format!(
                            "Cannot update from other segment because it is \
                             missing vector name {vector_name}; \
                             retry after optimizer config refresh"
                        ))
                    })?;

                    vector_data
                        .old_indices
                        .push(Arc::clone(&other_vector_data.vector_index));

                    Ok(other_vector_data.vector_storage.borrow())
                })
                .collect::<Result<Vec<_>, OperationError>>()?;

            let source_refs: Vec<&VectorStorageEnum> =
                other_vector_storages.iter().map(|s| &**s).collect();
            let internal_range = merge_from(
                &mut vector_data.vector_storage,
                &points_to_insert,
                &source_refs,
                stopped,
            )?;

            if new_internal_range != internal_range {
                debug_assert!(
                    new_internal_range != internal_range,
                    "Internal ids range mismatch between self segment vectors and other segment vectors\n\
                        vector_name: {vector_name}, self range: {new_internal_range:?}, other range: {internal_range:?}"
                );
                return Err(OperationError::service_error(format!(
                    "Internal ids range mismatch between self segment vectors and other segment vectors\n\
                        vector_name: {vector_name}, self range: {new_internal_range:?}, other range: {internal_range:?}"
                )));
            }
        }

        let hw_counter = HardwareCounterCell::disposable(); // Disposable counter for internal operations.

        let internal_id_iter = new_internal_range.zip(points_to_insert.iter());

        for (new_internal_id, point_data) in internal_id_iter {
            check_process_stopped(stopped)?;

            let old_internal_id = point_data.internal_id;

            let other_payload = payloads[point_data.segment_index.get() as usize]
                .with_view(|v| v.get_payload_sequential(old_internal_id, &hw_counter))?; // Internal operation, no measurement needed!

            match self.id_tracker.internal_id_with_behavior(
                ExtendedPointId::from(point_data.external_id),
                // Dedup guard on a freshly built target (no deferred heads
                // here): match any existing copy of this external id.
                DeferredBehavior::WithDeferred,
            ) {
                Some(existing_internal_id) => {
                    debug_assert!(
                        false,
                        "This code should not be reachable, cause points were resolved with `merged_points`"
                    );

                    let existing_external_version = self
                        .id_tracker
                        .internal_version(existing_internal_id)
                        .unwrap();

                    let remove_id = if existing_external_version < point_data.version {
                        // Other version is the newest, remove the existing one and replace
                        self.id_tracker
                            .drop(ExtendedPointId::from(point_data.external_id))?;
                        self.id_tracker.set_link(
                            ExtendedPointId::from(point_data.external_id),
                            new_internal_id,
                        )?;
                        self.id_tracker
                            .set_internal_version(new_internal_id, point_data.version)?;
                        self.payload_storage
                            .clear(existing_internal_id, &hw_counter)?;

                        existing_internal_id
                    } else {
                        // Old version is still good, do not move anything else
                        // Mark newly added vector as removed
                        new_internal_id
                    };
                    for vector_data in self.vector_data.values_mut() {
                        vector_data.vector_storage.delete_vector(remove_id)?;
                    }
                }
                None => {
                    self.id_tracker.set_link(
                        ExtendedPointId::from(point_data.external_id),
                        new_internal_id,
                    )?;
                    self.id_tracker
                        .set_internal_version(new_internal_id, point_data.version)?;
                }
            }

            // Propagate payload to new segment
            if !other_payload.is_empty() {
                self.payload_storage.set(
                    new_internal_id,
                    &other_payload,
                    &HardwareCounterCell::disposable(),
                )?;
            }
        }

        for payload in payloads {
            for (field, payload_schema) in payload.with_view(|v| v.indexed_fields()) {
                self.indexed_fields.insert(field, payload_schema);
            }
        }

        Ok(true)
    }

    /// Test wrapper for [`SegmentBuilder::build`].
    #[cfg(feature = "testing")]
    pub fn build_for_test(self, segments_path: &Path) -> Segment {
        use crate::segment::index::hnsw_index::get_num_indexing_threads;

        self.build(
            segments_path,
            Uuid::new_v4(),
            None,
            ResourcePermit::dummy(get_num_indexing_threads(0) as u32),
            &AtomicBool::new(false),
            &mut rand::rng(),
            &HardwareCounterCell::new(),
            ProgressTracker::new_for_test(),
        )
        .unwrap()
    }

    #[allow(clippy::too_many_arguments)]
    pub fn build<R: Rng + ?Sized>(
        self,
        segments_path: &Path,
        segment_uuid: Uuid,
        deferred_internal_id: Option<PointOffsetType>,
        permit: ResourcePermit,
        stopped: &AtomicBool,
        rng: &mut R,
        hw_counter: &HardwareCounterCell,
        progress_segment: ProgressTracker,
    ) -> Result<Segment, OperationError> {
        let temp_dir = {
            let SegmentBuilder {
                version,
                id_tracker,
                payload_storage,
                mut vector_data,
                segment_config,
                hnsw_global_config,
                temp_dir,
                indexed_fields,
                defragment_keys: _,
                live_vector_names: _,
            } = self;

            let progress_quantization = progress_segment.subtask("quantization");
            let progress_payload_index = progress_segment.subtask("payload_index");
            let indexed_fields = indexed_fields
                .into_iter()
                .map(|(field, payload_schema)| {
                    let progress = progress_payload_index
                        .subtask(format!("{}:{field}", payload_schema.name()));
                    (field, payload_schema, progress)
                })
                .collect::<Vec<(PayloadKeyType, PayloadFieldSchema, ProgressTracker)>>();
            let progress_vector_index = progress_segment.subtask("vector_index");
            let progress_sparse_vector_index = progress_segment.subtask("sparse_vector_index");

            let appendable_flag = segment_config.is_appendable();

            payload_storage.flusher()()?;
            let payload_storage_arc = Arc::new(AtomicRefCell::new(payload_storage));

            let id_tracker = match id_tracker {
                IdTrackerEnum::InMemoryIdTracker(in_memory_id_tracker) => {
                    // Serverless-compatible builds produce the disk-resident tracker
                    // (mapping stays on disk); otherwise the in-RAM immutable tracker.
                    if feature_flags().serverless_compatible() {
                        let disk_id_tracker = DiskIdTracker::from_in_memory_tracker(
                            &MmapFs,
                            in_memory_id_tracker,
                            temp_dir.path(),
                        )?;
                        IdTrackerEnum::DiskIdTracker(disk_id_tracker)
                    } else {
                        let (versions, mappings) = in_memory_id_tracker.into_internal();
                        let compressed_mapping = CompressedPointMappings::from_mappings(mappings);
                        let immutable_id_tracker = ImmutableIdTracker::new(
                            &MmapFs,
                            temp_dir.path(),
                            &versions,
                            compressed_mapping,
                        )?;
                        IdTrackerEnum::ImmutableIdTracker(immutable_id_tracker)
                    }
                }
                IdTrackerEnum::MutableIdTracker(_) => id_tracker,
                IdTrackerEnum::ImmutableIdTracker(_) => {
                    unreachable!("ImmutableIdTracker should not be used for building segment")
                }
                IdTrackerEnum::DiskIdTracker(_) => {
                    unreachable!("DiskIdTracker should not be used for building segment")
                }
            };

            id_tracker.mapping_flusher()()?;
            id_tracker.versions_flusher()()?;
            let id_tracker_arc = Arc::new(AtomicRefCell::new(id_tracker));

            let mut quantized_vectors = Self::update_quantization(
                &segment_config,
                &vector_data,
                temp_dir.path(),
                &permit,
                stopped,
                progress_quantization,
            )?;

            let mut vector_storages_arc = HashMap::new();
            let mut old_indices = HashMap::new();

            for vector_name in segment_config.vector_data.keys() {
                let Some(vector_info) = vector_data.remove(vector_name) else {
                    return Err(OperationError::service_error(format!(
                        "Vector storage for vector name {vector_name} not found on segment build"
                    )));
                };

                vector_info.vector_storage.flusher()()?;

                let vector_storage_arc = Arc::new(AtomicRefCell::new(vector_info.vector_storage));

                old_indices.insert(vector_name, vector_info.old_indices);

                vector_storages_arc.insert(vector_name.to_owned(), vector_storage_arc);
            }

            for vector_name in segment_config.sparse_vector_data.keys() {
                let Some(vector_info) = vector_data.remove(vector_name) else {
                    return Err(OperationError::service_error(format!(
                        "Vector storage for vector name {vector_name} not found on sparse segment build"
                    )));
                };

                vector_info.vector_storage.flusher()()?;

                let vector_storage_arc = Arc::new(AtomicRefCell::new(vector_info.vector_storage));

                vector_storages_arc.insert(vector_name.to_owned(), vector_storage_arc);
            }

            let payload_index_path = get_payload_index_path(temp_dir.path());

            progress_payload_index.start();
            let mut payload_index = StructPayloadIndex::open(
                payload_storage_arc.clone(),
                id_tracker_arc.clone(),
                vector_storages_arc.clone(),
                &payload_index_path,
                StorageType::from_appendable(appendable_flag),
                IndexLoadMode::CreateIfMissing,
            )?;
            for (field, payload_schema, progress) in indexed_fields {
                progress.start();
                payload_index.set_indexed(&field, payload_schema, hw_counter)?;
                check_process_stopped(stopped)?;
            }
            drop(progress_payload_index);

            payload_index.flusher()()?;
            let payload_index_arc = Arc::new(AtomicRefCell::new(payload_index));

            // Try to lock GPU device.
            #[cfg(feature = "gpu")]
            let gpu_devices_manager = crate::segment::index::hnsw_index::gpu::GPU_DEVICES_MANAGER.read();
            #[cfg(feature = "gpu")]
            let gpu_device = gpu_devices_manager
                .as_ref()
                .map(|devices_manager| devices_manager.lock_device(stopped))
                .transpose()?
                .flatten();
            #[cfg(not(feature = "gpu"))]
            let gpu_device = None;

            // Arc permit to share it with each vector store
            let permit = Arc::new(permit);

            progress_vector_index.start();
            for (vector_name, vector_config) in &segment_config.vector_data {
                let vector_storage = vector_storages_arc.remove(vector_name).unwrap();
                let quantized_vectors =
                    Arc::new(AtomicRefCell::new(quantized_vectors.remove(vector_name)));

                let index = build_vector_index(
                    vector_config,
                    VectorIndexOpenArgs {
                        path: &get_vector_index_path(temp_dir.path(), vector_name),
                        id_tracker: id_tracker_arc.clone(),
                        vector_storage: vector_storage.clone(),
                        payload_index: payload_index_arc.clone(),
                        quantized_vectors: quantized_vectors.clone(),
                    },
                    VectorIndexBuildArgs {
                        permit: permit.clone(),
                        old_indices: &old_indices.remove(vector_name).unwrap(),
                        gpu_device: gpu_device.as_ref(),
                        stopped,
                        rng,
                        hnsw_global_config: &hnsw_global_config,
                        feature_flags: feature_flags(),
                        progress: progress_vector_index.running_subtask(vector_name),
                    },
                )?;

                if vector_storage.borrow().is_on_disk() {
                    // If vector storage is expected to be on-disk, we need to clear cache
                    // to avoid cache pollution
                    vector_storage.borrow().clear_cache()?;
                }

                if let Some(quantized_vectors) = quantized_vectors.borrow().as_ref() {
                    quantized_vectors.clear_cache()?;
                }

                // Index if always loaded on-disk=true from build function
                // So we may clear unconditionally
                index.clear_cache()?;
            }
            drop(progress_vector_index);

            progress_sparse_vector_index.start();
            for (vector_name, sparse_vector_config) in &segment_config.sparse_vector_data {
                let vector_index_path = get_vector_index_path(temp_dir.path(), vector_name);

                let vector_storage_arc = vector_storages_arc.remove(vector_name).unwrap();

                let index = open_or_create_sparse_vector_index(SparseVectorIndexOpenArgs {
                    fs: &MmapFs,
                    config: sparse_vector_config.index,
                    id_tracker: id_tracker_arc.clone(),
                    vector_storage: vector_storage_arc.clone(),
                    payload_index: payload_index_arc.clone(),
                    path: &vector_index_path,
                    stopped,
                    tick_progress: || (),
                })?;

                if sparse_vector_config.storage_type.is_on_disk() {
                    // If vector storage is expected to be on-disk, we need to clear cache
                    // to avoid cache pollution
                    vector_storage_arc.borrow().clear_cache()?;
                }

                if sparse_vector_config.index.index_type.is_on_disk() {
                    index.clear_cache()?;
                }

                // Ensure we don't use the sparse vector index in future because it's missing proper setup of deferred points.
                drop(index);
            }
            drop(progress_sparse_vector_index);

            if segment_config.payload_storage_type.is_on_disk() {
                // If payload storage is expected to be on-disk, we need to clear cache
                // to avoid cache pollution
                payload_storage_arc.borrow().clear_cache()?;
            }

            // Clear cache for payload index to avoid cache pollution
            payload_index_arc.borrow().clear_cache_if_on_disk()?;

            // The id tracker is loaded into RAM but its on-disk files are written
            // during the build; drop their page cache to avoid cache pollution.
            id_tracker_arc.borrow().clear_cache_if_on_disk()?;

            // We're done with CPU-intensive tasks, release CPU permit
            debug_assert_eq!(
                Arc::strong_count(&permit),
                1,
                "Must release CPU permit Arc everywhere",
            );
            drop(permit);

            // Finalize the newly created segment by saving config and version
            Segment::save_state(
                &SegmentState {
                    initial_version: Some(version), // TODO!?
                    version: Some(version),
                    config: segment_config,
                },
                temp_dir.path(),
            )?;

            // After version is saved, segment can be loaded on restart
            SegmentVersion::save(temp_dir.path())?;
            // All temp data is evicted from RAM
            temp_dir
        };

        // Move fully constructed segment into collection directory and load back to RAM
        let destination_path = segments_path.join(segment_uuid.to_string());
        fs::rename(temp_dir.keep(), &destination_path)
            .describe("Moving segment data after optimization")?;

        load_segment(
            &destination_path,
            segment_uuid,
            deferred_internal_id,
            stopped,
        )
    }

    fn update_quantization(
        segment_config: &SegmentConfig,
        vector_storages: &HashMap<VectorNameBuf, VectorData>,
        temp_path: &Path,
        permit: &ResourcePermit,
        stopped: &AtomicBool,
        progress: ProgressTracker,
    ) -> OperationResult<HashMap<VectorNameBuf, QuantizedVectors>> {
        progress.start();
        let config = segment_config.clone();

        let mut quantized_vectors_map = HashMap::new();

        for (vector_name, vector_info) in vector_storages {
            let Some(vector_config) = config.vector_data.get(vector_name) else {
                continue;
            };

            let is_appendable = vector_config.is_appendable();

            // If appendable quantization feature is not enabled, skip appendable case.
            if is_appendable && !crate::common::flags::feature_flags().appendable_quantization {
                continue;
            }

            let max_threads = permit.num_cpus as usize;

            if let Some(quantization_config) = config.quantization_config(vector_name) {
                // Don't build quantization for appendable vectors if quantization method does not support it
                if is_appendable && !quantization_config.supports_appendable() {
                    continue;
                }

                let progress_vector = progress.running_subtask(vector_name);

                let segment_path = temp_path;
                let quantized_storage_type = if is_appendable {
                    QuantizedVectorsStorageType::Mutable
                } else {
                    QuantizedVectorsStorageType::Immutable
                };

                let vector_storage_path = get_vector_storage_path(segment_path, vector_name);

                let quantized_vectors = QuantizedVectors::create(
                    &vector_info.vector_storage,
                    quantization_config,
                    quantized_storage_type,
                    &vector_storage_path,
                    max_threads,
                    stopped,
                )?;

                quantized_vectors_map.insert(vector_name.to_owned(), quantized_vectors);

                drop(progress_vector);
            }
        }
        Ok(quantized_vectors_map)
    }

    /// Populate cache of all vector storages, so it will be faster to build index
    pub fn populate_vector_storages(&self) -> OperationResult<()> {
        for vector_data in self.vector_data.values() {
            vector_data.vector_storage.populate()?;
        }
        Ok(())
    }
}

fn uuid_hash<I>(hash: &mut u64, ids: I)
where
    I: Iterator<Item = u128>,
{
    for id in ids {
        let uuid = Uuid::from_u128(id);

        // Not all Uuid versions hold timestamp data. The most common version, v4 for example is completely
        // random and can't be sorted. To still allow defragmentation, we assume that usually the same
        // version gets used for a payload key and implement an alternative sorting criteria, that just
        // takes the Uuids bytes to group equal Uuids together.
        if let Some(timestamp) = uuid.get_timestamp() {
            *hash = hash.wrapping_add(timestamp.to_gregorian().0);
        } else {
            // First part of u128
            *hash = hash.wrapping_add((id >> 64) as u64);

            // Second part of u128
            *hash = hash.wrapping_add(id as u64);
        }
    }
}

fn create_temp_dir(parent_path: &Path) -> Result<TempDir, OperationError> {
    // Ensure parent path exists
    fs::create_dir_all(parent_path)
        .and_then(|_| TempDir::with_prefix_in("segment_builder_", parent_path))
        .map_err(|err| {
            OperationError::service_error(format!(
                "Could not create temp directory in `{}`: {}",
                parent_path.display(),
                err
            ))
        })
}