core/src/read_options.rs - paimon-vector-index - Git at Google

 // Licensed to the Apache Software Foundation (ASF) under one
 // or more contributor license agreements.  See the NOTICE file
 // distributed with this work for additional information
 // regarding copyright ownership.  The ASF licenses this file
 // to you under the Apache License, Version 2.0 (the
 // "License"); you may not use this file except in compliance
 // with the License.  You may obtain a copy of the License at
 //
 //   http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing,
 // software distributed under the License is distributed on an
 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.

 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
 /// Build-time serving intent used to choose DiskANN's persisted layout.
 ///
 /// Reader I/O policy is always derived from input capabilities and measured
 /// latency; this enum is not a Reader option.
 pub enum DeploymentProfile {
     #[default]
     Auto = 0,
     Memory = 1,
     LocalStorage = 2,
     RemoteStorage = 3,
     ObjectStore = 4,
 }

 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(crate) struct ReadPlan {
     pub window_bytes: usize,
     pub graph_beam_width: usize,
     pub filtered_graph_beam_width: usize,
 }

 impl DeploymentProfile {
     pub(crate) const fn read_plan(self) -> ReadPlan {
         match self {
             Self::Auto | Self::LocalStorage => ReadPlan {
                 window_bytes: 16 * 1024,
                 graph_beam_width: 4,
                 filtered_graph_beam_width: 4,
             },
             Self::Memory => ReadPlan {
                 window_bytes: 4 * 1024,
                 graph_beam_width: 16,
                 filtered_graph_beam_width: 4,
             },
             Self::RemoteStorage => ReadPlan {
                 window_bytes: 32 * 1024,
                 graph_beam_width: 16,
                 filtered_graph_beam_width: 4,
             },
             Self::ObjectStore => ReadPlan {
                 window_bytes: 64 * 1024,
                 graph_beam_width: 16,
                 filtered_graph_beam_width: 4,
             },
         }
     }
 }

 impl ReadPlan {
     pub(crate) fn with_capabilities(
         mut self,
         capabilities: crate::io::SeekReadCapabilities,
     ) -> Self {
         if capabilities.preferred_window_bytes > 0 {
             self.window_bytes = capabilities.preferred_window_bytes;
         }
         // Bound pathological adapter hints while keeping enough room for one
         // complete DiskANN page. The planner clips the final window to the
         // section, so a non-page-sized capability hint must not split a page
         // across two windows.
         self.window_bytes = self.window_bytes.clamp(4 * 1024, 1024 * 1024);
         self.window_bytes = self.window_bytes.div_ceil(4 * 1024) * (4 * 1024);
         self.window_bytes = self.window_bytes.min(1024 * 1024);
         if capabilities.max_ranges_per_pread > 0 {
             self.graph_beam_width = self
                 .graph_beam_width
                 .min(capabilities.max_ranges_per_pread)
                 .max(1);
             self.filtered_graph_beam_width = self
                 .filtered_graph_beam_width
                 .min(capabilities.max_ranges_per_pread)
                 .max(1);
         }
         self
     }
 }

 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct VectorIndexReaderOptions {
     /// Total bytes available to required resident state and retained caches.
     pub memory_budget_bytes: usize,
     cache_overrides: Option<ResolvedVectorIndexReaderOptions>,
 }

 /// The concrete DiskANN read policy derived from input capabilities, observed
 /// latency, index layout, and the Reader memory budget.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct VectorIndexReadPlan {
     pub random_read_latency_nanos: u64,
     pub window_bytes: usize,
     pub max_ranges_per_read: usize,
     pub graph_beam_width: usize,
     pub filtered_graph_beam_width: usize,
     pub adjacency_preload_bytes: usize,
     pub adjacency_cache_bytes: usize,
     pub raw_vector_cache_bytes: usize,
     pub memory_budget_bytes: usize,
 }

 impl Default for VectorIndexReaderOptions {
     fn default() -> Self {
         Self {
             memory_budget_bytes: 4 * 1024 * 1024 * 1024,
             cache_overrides: None,
         }
     }
 }

 impl VectorIndexReaderOptions {
     pub fn new(memory_budget_bytes: usize) -> Self {
         Self {
             memory_budget_bytes,
             cache_overrides: None,
         }
     }

     pub(crate) fn resolve_cache_budgets(
         self,
         read_tier: DeploymentProfile,
         resident_steady_bytes: usize,
         adjacency_section_bytes: usize,
         raw_vector_section_bytes: usize,
     ) -> ResolvedVectorIndexReaderOptions {
         if let Some(mut overrides) = self.cache_overrides {
             overrides.read_tier = read_tier;
             return overrides;
         }
         let available = self
             .memory_budget_bytes
             .saturating_sub(resident_steady_bytes);
         let preload_cap = match read_tier {
             DeploymentProfile::Auto
             | DeploymentProfile::Memory
             | DeploymentProfile::LocalStorage => 16 * 1024 * 1024,
             DeploymentProfile::RemoteStorage => 32 * 1024 * 1024,
             DeploymentProfile::ObjectStore => 64 * 1024 * 1024,
         };
         let adjacency_preload_bytes = (available / 2)
             .min(preload_cap)
             .min(adjacency_section_bytes);
         let after_preload = available.saturating_sub(adjacency_preload_bytes);
         let per_cache_cap = match read_tier {
             DeploymentProfile::RemoteStorage | DeploymentProfile::ObjectStore => usize::MAX,
             DeploymentProfile::Auto
             | DeploymentProfile::Memory
             | DeploymentProfile::LocalStorage => 64 * 1024 * 1024,
         };
         let cold_adjacency_bytes = adjacency_section_bytes.saturating_sub(adjacency_preload_bytes);
         let mut adjacency_cache_bytes = (after_preload / 2)
             .min(per_cache_cap)
             .min(cold_adjacency_bytes);
         let mut raw_vector_cache_bytes = after_preload
             .saturating_sub(adjacency_cache_bytes)
             .min(per_cache_cap)
             .min(raw_vector_section_bytes);
         let mut unassigned = after_preload
             .saturating_sub(adjacency_cache_bytes)
             .saturating_sub(raw_vector_cache_bytes);
         let adjacency_extra = unassigned
             .min(per_cache_cap.saturating_sub(adjacency_cache_bytes))
             .min(cold_adjacency_bytes.saturating_sub(adjacency_cache_bytes));
         adjacency_cache_bytes = adjacency_cache_bytes.saturating_add(adjacency_extra);
         unassigned = unassigned.saturating_sub(adjacency_extra);
         let vector_extra = unassigned
             .min(per_cache_cap.saturating_sub(raw_vector_cache_bytes))
             .min(raw_vector_section_bytes.saturating_sub(raw_vector_cache_bytes));
         raw_vector_cache_bytes = raw_vector_cache_bytes.saturating_add(vector_extra);
         ResolvedVectorIndexReaderOptions {
             read_tier,
             adjacency_preload_bytes,
             adjacency_cache_bytes,
             max_resident_bytes: self.memory_budget_bytes,
             raw_vector_cache_bytes,
         }
     }

     #[cfg(test)]
     pub(crate) fn with_cache_budgets(
         read_tier: DeploymentProfile,
         adjacency_preload_bytes: usize,
         adjacency_cache_bytes: usize,
         max_resident_bytes: usize,
         raw_vector_cache_bytes: usize,
     ) -> Self {
         Self {
             memory_budget_bytes: max_resident_bytes,
             cache_overrides: Some(ResolvedVectorIndexReaderOptions {
                 read_tier,
                 adjacency_preload_bytes,
                 adjacency_cache_bytes,
                 max_resident_bytes,
                 raw_vector_cache_bytes,
             }),
         }
     }
 }

 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(crate) struct ResolvedVectorIndexReaderOptions {
     pub read_tier: DeploymentProfile,
     pub adjacency_preload_bytes: usize,
     pub adjacency_cache_bytes: usize,
     pub max_resident_bytes: usize,
     pub raw_vector_cache_bytes: usize,
 }

 #[cfg(test)]
 mod tests {
     use crate::io::SeekReadCapabilities;

     use super::{DeploymentProfile, VectorIndexReaderOptions};

     #[test]
     fn deployment_profiles_have_distinct_read_plans() {
         let options = VectorIndexReaderOptions::default();
         assert_eq!(options.memory_budget_bytes, 4 * 1024 * 1024 * 1024);

         let resolved = options.resolve_cache_budgets(
             DeploymentProfile::Auto,
             1024,
             16 * 1024 * 1024,
             8 * 1024 * 1024,
         );
         assert!(resolved.adjacency_preload_bytes > 0);
         assert_eq!(
             resolved.adjacency_preload_bytes + resolved.adjacency_cache_bytes,
             16 * 1024 * 1024
         );
         assert!(resolved.raw_vector_cache_bytes > 0);
         assert!(
             1024 + resolved.adjacency_preload_bytes
                 + resolved.adjacency_cache_bytes
                 + resolved.raw_vector_cache_bytes
                 <= options.memory_budget_bytes
         );

         let memory = DeploymentProfile::Memory.read_plan();
         let local = DeploymentProfile::LocalStorage.read_plan();
         let remote = DeploymentProfile::RemoteStorage.read_plan();
         let object_store = DeploymentProfile::ObjectStore.read_plan();
         assert_eq!(
             (
                 memory.window_bytes,
                 memory.graph_beam_width,
                 memory.filtered_graph_beam_width
             ),
             (4096, 16, 4)
         );
         assert_eq!(
             (
                 local.window_bytes,
                 local.graph_beam_width,
                 local.filtered_graph_beam_width
             ),
             (16 * 1024, 4, 4)
         );
         assert_eq!(
             (
                 remote.window_bytes,
                 remote.graph_beam_width,
                 remote.filtered_graph_beam_width
             ),
             (32 * 1024, 16, 4)
         );
         assert_eq!(
             (
                 object_store.window_bytes,
                 object_store.graph_beam_width,
                 object_store.filtered_graph_beam_width
             ),
             (64 * 1024, 16, 4)
         );
     }

     #[test]
     fn remote_cache_budgets_use_available_memory_beyond_the_local_cache_caps() {
         const MIB: usize = 1024 * 1024;
         let resolved = VectorIndexReaderOptions::new(4 * 1024 * MIB).resolve_cache_budgets(
             DeploymentProfile::RemoteStorage,
             256 * MIB,
             512 * MIB,
             2 * 1024 * MIB,
         );

         assert_eq!(resolved.adjacency_preload_bytes, 32 * MIB);
         assert_eq!(
             resolved.adjacency_cache_bytes,
             512 * MIB - resolved.adjacency_preload_bytes
         );
         assert_eq!(resolved.raw_vector_cache_bytes, 2 * 1024 * MIB);
         assert!(
             256 * MIB
                 + resolved.adjacency_preload_bytes
                 + resolved.adjacency_cache_bytes
                 + resolved.raw_vector_cache_bytes
                 <= resolved.max_resident_bytes
         );
     }

     #[test]
     fn capability_windows_are_bounded_and_keep_complete_diskann_pages() {
         let plan = DeploymentProfile::LocalStorage
             .read_plan()
             .with_capabilities(SeekReadCapabilities {
                 estimated_random_read_latency_nanos: 0,
                 preferred_window_bytes: 10_000,
                 max_ranges_per_pread: 2,
             });
         assert_eq!(plan.window_bytes, 12 * 1024);
         assert_eq!(plan.graph_beam_width, 2);
         assert_eq!(plan.filtered_graph_beam_width, 2);

         let bounded = DeploymentProfile::RemoteStorage
             .read_plan()
             .with_capabilities(SeekReadCapabilities {
                 estimated_random_read_latency_nanos: 0,
                 preferred_window_bytes: usize::MAX,
                 max_ranges_per_pread: 0,
             });
         assert_eq!(bounded.window_bytes, 1024 * 1024);
     }
 }
	// Licensed to the Apache Software Foundation (ASF) under one
	// or more contributor license agreements. See the NOTICE file
	// distributed with this work for additional information
	// regarding copyright ownership. The ASF licenses this file
	// to you under the Apache License, Version 2.0 (the
	// "License"); you may not use this file except in compliance
	// with the License. You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing,
	// software distributed under the License is distributed on an
	// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	// KIND, either express or implied. See the License for the
	// specific language governing permissions and limitations
	// under the License.

	#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
	/// Build-time serving intent used to choose DiskANN's persisted layout.
	///
	/// Reader I/O policy is always derived from input capabilities and measured
	/// latency; this enum is not a Reader option.
	pub enum DeploymentProfile {
	#[default]
	Auto = 0,
	Memory = 1,
	LocalStorage = 2,
	RemoteStorage = 3,
	ObjectStore = 4,
	}

	#[derive(Debug, Clone, Copy, PartialEq, Eq)]
	pub(crate) struct ReadPlan {
	pub window_bytes: usize,
	pub graph_beam_width: usize,
	pub filtered_graph_beam_width: usize,
	}

	impl DeploymentProfile {
	pub(crate) const fn read_plan(self) -> ReadPlan {
	match self {
	Self::Auto \| Self::LocalStorage => ReadPlan {
	window_bytes: 16 * 1024,
	graph_beam_width: 4,
	filtered_graph_beam_width: 4,
	},
	Self::Memory => ReadPlan {
	window_bytes: 4 * 1024,
	graph_beam_width: 16,
	filtered_graph_beam_width: 4,
	},
	Self::RemoteStorage => ReadPlan {
	window_bytes: 32 * 1024,
	graph_beam_width: 16,
	filtered_graph_beam_width: 4,
	},
	Self::ObjectStore => ReadPlan {
	window_bytes: 64 * 1024,
	graph_beam_width: 16,
	filtered_graph_beam_width: 4,
	},
	}
	}
	}

	impl ReadPlan {
	pub(crate) fn with_capabilities(
	mut self,
	capabilities: crate::io::SeekReadCapabilities,
	) -> Self {
	if capabilities.preferred_window_bytes > 0 {
	self.window_bytes = capabilities.preferred_window_bytes;
	}
	// Bound pathological adapter hints while keeping enough room for one
	// complete DiskANN page. The planner clips the final window to the
	// section, so a non-page-sized capability hint must not split a page
	// across two windows.
	self.window_bytes = self.window_bytes.clamp(4 * 1024, 1024 * 1024);
	self.window_bytes = self.window_bytes.div_ceil(4 * 1024) * (4 * 1024);
	self.window_bytes = self.window_bytes.min(1024 * 1024);
	if capabilities.max_ranges_per_pread > 0 {
	self.graph_beam_width = self
	.graph_beam_width
	.min(capabilities.max_ranges_per_pread)
	.max(1);
	self.filtered_graph_beam_width = self
	.filtered_graph_beam_width
	.min(capabilities.max_ranges_per_pread)
	.max(1);
	}
	self
	}
	}

	#[derive(Debug, Clone, Copy, PartialEq, Eq)]
	pub struct VectorIndexReaderOptions {
	/// Total bytes available to required resident state and retained caches.
	pub memory_budget_bytes: usize,
	cache_overrides: Option<ResolvedVectorIndexReaderOptions>,
	}

	/// The concrete DiskANN read policy derived from input capabilities, observed
	/// latency, index layout, and the Reader memory budget.
	#[derive(Debug, Clone, Copy, PartialEq, Eq)]
	pub struct VectorIndexReadPlan {
	pub random_read_latency_nanos: u64,
	pub window_bytes: usize,
	pub max_ranges_per_read: usize,
	pub graph_beam_width: usize,
	pub filtered_graph_beam_width: usize,
	pub adjacency_preload_bytes: usize,
	pub adjacency_cache_bytes: usize,
	pub raw_vector_cache_bytes: usize,
	pub memory_budget_bytes: usize,
	}

	impl Default for VectorIndexReaderOptions {
	fn default() -> Self {
	Self {
	memory_budget_bytes: 4 * 1024 * 1024 * 1024,
	cache_overrides: None,
	}
	}
	}

	impl VectorIndexReaderOptions {
	pub fn new(memory_budget_bytes: usize) -> Self {
	Self {
	memory_budget_bytes,
	cache_overrides: None,
	}
	}

	pub(crate) fn resolve_cache_budgets(
	self,
	read_tier: DeploymentProfile,
	resident_steady_bytes: usize,
	adjacency_section_bytes: usize,
	raw_vector_section_bytes: usize,
	) -> ResolvedVectorIndexReaderOptions {
	if let Some(mut overrides) = self.cache_overrides {
	overrides.read_tier = read_tier;
	return overrides;
	}
	let available = self
	.memory_budget_bytes
	.saturating_sub(resident_steady_bytes);
	let preload_cap = match read_tier {
	DeploymentProfile::Auto
	\| DeploymentProfile::Memory
	\| DeploymentProfile::LocalStorage => 16 * 1024 * 1024,
	DeploymentProfile::RemoteStorage => 32 * 1024 * 1024,
	DeploymentProfile::ObjectStore => 64 * 1024 * 1024,
	};
	let adjacency_preload_bytes = (available / 2)
	.min(preload_cap)
	.min(adjacency_section_bytes);
	let after_preload = available.saturating_sub(adjacency_preload_bytes);
	let per_cache_cap = match read_tier {
	DeploymentProfile::RemoteStorage \| DeploymentProfile::ObjectStore => usize::MAX,
	DeploymentProfile::Auto
	\| DeploymentProfile::Memory
	\| DeploymentProfile::LocalStorage => 64 * 1024 * 1024,
	};
	let cold_adjacency_bytes = adjacency_section_bytes.saturating_sub(adjacency_preload_bytes);
	let mut adjacency_cache_bytes = (after_preload / 2)
	.min(per_cache_cap)
	.min(cold_adjacency_bytes);
	let mut raw_vector_cache_bytes = after_preload
	.saturating_sub(adjacency_cache_bytes)
	.min(per_cache_cap)
	.min(raw_vector_section_bytes);
	let mut unassigned = after_preload
	.saturating_sub(adjacency_cache_bytes)
	.saturating_sub(raw_vector_cache_bytes);
	let adjacency_extra = unassigned
	.min(per_cache_cap.saturating_sub(adjacency_cache_bytes))
	.min(cold_adjacency_bytes.saturating_sub(adjacency_cache_bytes));
	adjacency_cache_bytes = adjacency_cache_bytes.saturating_add(adjacency_extra);
	unassigned = unassigned.saturating_sub(adjacency_extra);
	let vector_extra = unassigned
	.min(per_cache_cap.saturating_sub(raw_vector_cache_bytes))
	.min(raw_vector_section_bytes.saturating_sub(raw_vector_cache_bytes));
	raw_vector_cache_bytes = raw_vector_cache_bytes.saturating_add(vector_extra);
	ResolvedVectorIndexReaderOptions {
	read_tier,
	adjacency_preload_bytes,
	adjacency_cache_bytes,
	max_resident_bytes: self.memory_budget_bytes,
	raw_vector_cache_bytes,
	}
	}

	#[cfg(test)]
	pub(crate) fn with_cache_budgets(
	read_tier: DeploymentProfile,
	adjacency_preload_bytes: usize,
	adjacency_cache_bytes: usize,
	max_resident_bytes: usize,
	raw_vector_cache_bytes: usize,
	) -> Self {
	Self {
	memory_budget_bytes: max_resident_bytes,
	cache_overrides: Some(ResolvedVectorIndexReaderOptions {
	read_tier,
	adjacency_preload_bytes,
	adjacency_cache_bytes,
	max_resident_bytes,
	raw_vector_cache_bytes,
	}),
	}
	}
	}

	#[derive(Debug, Clone, Copy, PartialEq, Eq)]
	pub(crate) struct ResolvedVectorIndexReaderOptions {
	pub read_tier: DeploymentProfile,
	pub adjacency_preload_bytes: usize,
	pub adjacency_cache_bytes: usize,
	pub max_resident_bytes: usize,
	pub raw_vector_cache_bytes: usize,
	}

	#[cfg(test)]
	mod tests {
	use crate::io::SeekReadCapabilities;

	use super::{DeploymentProfile, VectorIndexReaderOptions};

	#[test]
	fn deployment_profiles_have_distinct_read_plans() {
	let options = VectorIndexReaderOptions::default();
	assert_eq!(options.memory_budget_bytes, 4 * 1024 * 1024 * 1024);

	let resolved = options.resolve_cache_budgets(
	DeploymentProfile::Auto,
	1024,
	16 * 1024 * 1024,
	8 * 1024 * 1024,
	);
	assert!(resolved.adjacency_preload_bytes > 0);
	assert_eq!(
	resolved.adjacency_preload_bytes + resolved.adjacency_cache_bytes,
	16 * 1024 * 1024
	);
	assert!(resolved.raw_vector_cache_bytes > 0);
	assert!(
	1024 + resolved.adjacency_preload_bytes
	+ resolved.adjacency_cache_bytes
	+ resolved.raw_vector_cache_bytes
	<= options.memory_budget_bytes
	);

	let memory = DeploymentProfile::Memory.read_plan();
	let local = DeploymentProfile::LocalStorage.read_plan();
	let remote = DeploymentProfile::RemoteStorage.read_plan();
	let object_store = DeploymentProfile::ObjectStore.read_plan();
	assert_eq!(
	(
	memory.window_bytes,
	memory.graph_beam_width,
	memory.filtered_graph_beam_width
	),
	(4096, 16, 4)
	);
	assert_eq!(
	(
	local.window_bytes,
	local.graph_beam_width,
	local.filtered_graph_beam_width
	),
	(16 * 1024, 4, 4)
	);
	assert_eq!(
	(
	remote.window_bytes,
	remote.graph_beam_width,
	remote.filtered_graph_beam_width
	),
	(32 * 1024, 16, 4)
	);
	assert_eq!(
	(
	object_store.window_bytes,
	object_store.graph_beam_width,
	object_store.filtered_graph_beam_width
	),
	(64 * 1024, 16, 4)
	);
	}

	#[test]
	fn remote_cache_budgets_use_available_memory_beyond_the_local_cache_caps() {
	const MIB: usize = 1024 * 1024;
	let resolved = VectorIndexReaderOptions::new(4 * 1024 * MIB).resolve_cache_budgets(
	DeploymentProfile::RemoteStorage,
	256 * MIB,
	512 * MIB,
	2 * 1024 * MIB,
	);

	assert_eq!(resolved.adjacency_preload_bytes, 32 * MIB);
	assert_eq!(
	resolved.adjacency_cache_bytes,
	512 * MIB - resolved.adjacency_preload_bytes
	);
	assert_eq!(resolved.raw_vector_cache_bytes, 2 * 1024 * MIB);
	assert!(
	256 * MIB
	+ resolved.adjacency_preload_bytes
	+ resolved.adjacency_cache_bytes
	+ resolved.raw_vector_cache_bytes
	<= resolved.max_resident_bytes
	);
	}

	#[test]
	fn capability_windows_are_bounded_and_keep_complete_diskann_pages() {
	let plan = DeploymentProfile::LocalStorage
	.read_plan()
	.with_capabilities(SeekReadCapabilities {
	estimated_random_read_latency_nanos: 0,
	preferred_window_bytes: 10_000,
	max_ranges_per_pread: 2,
	});
	assert_eq!(plan.window_bytes, 12 * 1024);
	assert_eq!(plan.graph_beam_width, 2);
	assert_eq!(plan.filtered_graph_beam_width, 2);

	let bounded = DeploymentProfile::RemoteStorage
	.read_plan()
	.with_capabilities(SeekReadCapabilities {
	estimated_random_read_latency_nanos: 0,
	preferred_window_bytes: usize::MAX,
	max_ranges_per_pread: 0,
	});
	assert_eq!(bounded.window_bytes, 1024 * 1024);
	}
	}