mod ah;
mod artifact;
mod binary;
mod config;
mod engine;
mod fast_scan;
mod geometry;
mod payload;
pub use ah::{
AhCodebook, AhEncodeScratch, AhQuery, CENTERS_PER_BLOCK, DEFAULT_ANISOTROPIC_THRESHOLD,
};
pub use artifact::{
SCANN_GLOBAL_ARTIFACT_VERSION, ScannAhCodebook, ScannAhCodebookRef, ScannRoutingLevel,
ScannRoutingLevelRef, ScannTrainedArtifact, ScannTrainedArtifactView,
};
pub use binary::{
BinaryScannHit, BinaryScannModel, BinaryScannProbePlan, BinaryScannSearchScratch,
BinaryScannSegment, BinaryScannSpillAssignment, BinaryScannTraining, BinaryScannTrainingStats,
QuantizedBinaryScannModel, QuantizedBinaryScannModelView,
};
pub use config::{
DEFAULT_TRAINING_SAMPLE_SIZE, MAX_SCANN_LEAVES, MAX_SCANN_TREE_LEVELS,
MIN_PARTITION_TRAINING_POINTS_PER_LEAF, MIN_POINTS_FOR_PARTITIONING,
PARTITION_TRAINING_POINTS_PER_CENTROID, SCANN_FAST_SCAN_LANES, ScannConfig, ScannEncoding,
ScannTrainingState,
};
pub use engine::{
EncodedFloatVector, FloatEncodeScratch, FloatRoutingTree, FloatScannModel, FloatScannQuery,
QuantizedFloatScannModel, QuantizedFloatScannModelView, RoutedLeaf, RoutingScratch,
RoutingTraining, RoutingTrainingStats, train_routing_tree,
};
pub use fast_scan::{FAST_SCAN_LANES, FastScanQuery, pack_fast_scan_block};
pub use geometry::{
QUALITY_OPTIMIZED_POINTS_PER_LEAF, ScannGeometry, derive_geometry, derive_geometry_with_levels,
derive_geometry_with_levels_and_sample_limit, desired_training_sample, geometry_for_leaves,
geometry_for_leaves_with_auto_depth,
};
pub use payload::{SCANN_SEGMENT_PAYLOAD_VERSION, ScannLeafRun, ScannSegmentPayload};
use std::fmt;
pub(crate) fn routing_prefix_for_child_coverage<T>(
ranked: &[T],
child_offsets: &[u32],
minimum_width: usize,
requested_children: usize,
node_id: impl Fn(&T) -> usize,
) -> usize {
let minimum_width = minimum_width.min(ranked.len());
let reachable = ranked.iter().fold(0usize, |total, item| {
let node = node_id(item);
total.saturating_add(
child_offsets[node + 1]
.saturating_sub(child_offsets[node])
.try_into()
.unwrap_or(usize::MAX),
)
});
let target = requested_children.min(reachable);
let mut covered = 0usize;
for (index, item) in ranked.iter().enumerate() {
let node = node_id(item);
covered = covered.saturating_add(
child_offsets[node + 1]
.saturating_sub(child_offsets[node])
.try_into()
.unwrap_or(usize::MAX),
);
let width = index + 1;
if width >= minimum_width && covered >= target {
return width;
}
}
ranked.len()
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ScannFormatError(String);
impl ScannFormatError {
pub(crate) fn new(message: impl Into<String>) -> Self {
Self(message.into())
}
}
impl fmt::Display for ScannFormatError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
}
}
impl std::error::Error for ScannFormatError {}
pub type ScannResult<T> = Result<T, ScannFormatError>;
#[cfg(test)]
mod tests {
use super::routing_prefix_for_child_coverage;
#[test]
fn routing_prefix_widens_until_requested_children_are_reachable() {
let ranked = [2usize, 0, 1];
let offsets = [0u32, 2, 5, 9];
assert_eq!(
routing_prefix_for_child_coverage(&ranked, &offsets, 1, 6, |node| *node),
2
);
assert_eq!(
routing_prefix_for_child_coverage(&ranked, &offsets, 1, 9, |node| *node),
3
);
assert_eq!(
routing_prefix_for_child_coverage(&ranked, &offsets, 2, 1, |node| *node),
2
);
}
}