upstream-archive/qdrant
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant.git synced 2026-09-21 13:37:46 -05:00
Code Issues Packages Projects Releases Wiki Activity

[UIO] Explicit padding for numeric point (#8668)

Browse Source 
* derive Pod for Point<T>

* remove clones on Copy type

* fix annoying clippy lint

* serde skip padding

* extract point into a new file

* assert aligment
This commit is contained in:
Luis Cossío
2026-05-08 13:46:46 +02:00
committed by timvisee
parent 94aa170527
commit c22a3f644c
11 changed files with 263 additions and 224 deletions
Download Patch File Download Diff File Expand all files Collapse all files
lib/segment/src/index/field_index/histogram.rs
+34 -174
View File
@@ -1,16 +1,17 @@
use std::collections::BTreeMap;
use std::collections::Bound::{Excluded, Included, Unbounded};
use std::fmt::Debug;
use std::ops::Bound;
use std::path::{Path, PathBuf};
use common::fs::{atomic_save_bin, atomic_save_json, read_bin, read_json};
use common::types::PointOffsetType;
use itertools::Itertools;
use num_traits::Num;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use crate::common::operation_error::OperationResult;
use crate::index::field_index::numeric_point::{Numericable, Point};
use crate::index::field_index::utils::check_boundaries;
const MIN_BUCKET_SIZE: usize = 10;
@@ -23,114 +24,6 @@ pub struct Counts {
pub right: usize,
}
#[allow(clippy::derive_ord_xor_partial_ord)]
#[derive(PartialEq, PartialOrd, Debug, Clone, Serialize, Deserialize)]
#[repr(C)]
pub struct Point<T> {
pub val: T,
pub idx: PointOffsetType,
}
impl<T> Point<T> {
pub fn new(val: T, idx: PointOffsetType) -> Self {
Self { val, idx }
}
}
impl<T: PartialEq> Eq for Point<T> {}
impl<T: PartialOrd + Copy> Ord for Point<T> {
fn cmp(&self, other: &Point<T>) -> std::cmp::Ordering {
(self.val, self.idx)
.partial_cmp(&(other.val, other.idx))
.unwrap()
}
}
/// A trait that should represent common properties of integer and floating point types.
/// In particular, i64 and f64.
pub trait Numericable: Num + PartialEq + PartialOrd + Copy {
fn min_value() -> Self;
fn max_value() -> Self;
fn to_f64(self) -> f64;
fn from_f64(x: f64) -> Self;
fn from_u128(x: u128) -> Self;
fn min(self, b: Self) -> Self {
if self < b { self } else { b }
}
fn max(self, b: Self) -> Self {
if self > b { self } else { b }
}
fn abs_diff(self, b: Self) -> Self {
if self > b { self - b } else { b - self }
}
}
impl Numericable for i64 {
fn min_value() -> Self {
i64::MIN
}
fn max_value() -> Self {
i64::MAX
}
fn to_f64(self) -> f64 {
self as f64
}
fn from_f64(x: f64) -> Self {
x as Self
}
fn from_u128(x: u128) -> Self {
x as i64
}
fn abs_diff(self, b: Self) -> Self {
i64::abs_diff(self, b) as i64
}
}
impl Numericable for f64 {
fn min_value() -> Self {
f64::MIN
}
fn max_value() -> Self {
f64::MAX
}
fn to_f64(self) -> f64 {
self
}
fn from_f64(x: f64) -> Self {
x
}
fn from_u128(x: u128) -> Self {
x as Self
}
}
impl Numericable for u128 {
fn min_value() -> Self {
u128::MIN
}
fn max_value() -> Self {
u128::MAX
}
fn to_f64(self) -> f64 {
self as f64
}
fn from_f64(x: f64) -> Self {
x as u128
}
fn from_u128(x: u128) -> Self {
x
}
fn abs_diff(self, b: Self) -> Self {
u128::abs_diff(self, b)
}
}
#[derive(Default, Debug, PartialEq)]
pub struct Histogram<T: Numericable + Serialize + DeserializeOwned> {
max_bucket_size: usize,
@@ -186,11 +79,8 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
},
)?;
let borders: Vec<(Point<T>, Counts)> = self
.borders
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
let borders: Vec<(Point<T>, Counts)> =
self.borders.iter().map(|(k, v)| (*k, v.clone())).collect();
atomic_save_bin(&borders_path, &borders)?;
Ok(())
}
@@ -229,14 +119,8 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
/// Returns `Unbounded` if there are no points stored
pub fn get_range_by_size(&self, from: Bound<T>, range_size: usize) -> Bound<T> {
let from_ = match from {
Included(val) => Included(Point {
val,
idx: PointOffsetType::MIN,
}),
Excluded(val) => Excluded(Point {
val,
idx: PointOffsetType::MAX,
}),
Included(val) => Included(Point::new(val, PointOffsetType::MIN)),
Excluded(val) => Excluded(Point::new(val, PointOffsetType::MAX)),
Unbounded => Unbounded,
};
@@ -256,26 +140,14 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
pub fn estimate(&self, from: Bound<T>, to: Bound<T>) -> (usize, usize, usize) {
let from_ = match &from {
Included(val) => Included(Point {
val: *val,
idx: PointOffsetType::MIN,
}),
Excluded(val) => Excluded(Point {
val: *val,
idx: PointOffsetType::MAX,
}),
Included(val) => Included(Point::new(*val, PointOffsetType::MIN)),
Excluded(val) => Excluded(Point::new(*val, PointOffsetType::MAX)),
Unbounded => Unbounded,
};
let to_ = match &to {
Included(val) => Included(Point {
val: *val,
idx: PointOffsetType::MAX,
}),
Excluded(val) => Excluded(Point {
val: *val,
idx: PointOffsetType::MIN,
}),
Included(val) => Included(Point::new(*val, PointOffsetType::MAX)),
Excluded(val) => Excluded(Point::new(*val, PointOffsetType::MIN)),
Unbounded => Unbounded,
};
@@ -296,7 +168,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
if matches!(from_, Unbounded) {
None
} else {
self.borders.range((Unbounded, from_.clone())).next_back()
self.borders.range((Unbounded, from_)).next_back()
}
};
@@ -304,7 +176,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
if matches!(to_, Unbounded) {
None
} else {
self.borders.range((to_.clone(), Unbounded)).next()
self.borders.range((to_, Unbounded)).next()
}
};
@@ -359,12 +231,12 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
let (mut close_neighbors, (mut far_left_neighbor, mut far_right_neighbor)) = {
let mut left_iterator = self
.borders
.range((Unbounded, Included(val.clone())))
.map(|(k, v)| (k.clone(), v.clone()));
.range((Unbounded, Included(val)))
.map(|(k, v)| (*k, v.clone()));
let mut right_iterator = self
.borders
.range((Excluded(val.clone()), Unbounded))
.map(|(k, v)| (k.clone(), v.clone()));
.range((Excluded(val), Unbounded))
.map(|(k, v)| (*k, v.clone()));
(
(left_iterator.next_back(), right_iterator.next()),
(left_iterator.next_back(), right_iterator.next()),
@@ -380,7 +252,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
if left_border_count.left == 0 {
// ...||
// ...|
(Some(left_border.clone()), None, true)
(Some(*left_border), None, true)
} else {
// ...|..|
// ...|.|
@@ -395,7 +267,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
},
);
(
Some(left_border.clone()),
Some(*left_border),
Some((new_border, new_border_count)),
true,
)
@@ -411,7 +283,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
if right_border_count.right == 0 {
// ||...
// |...
(Some(right_border.clone()), None, true)
(Some(*right_border), None, true)
} else {
// |..|...
// |.|...
@@ -426,7 +298,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
},
);
(
Some(right_border.clone()),
Some(*right_border),
Some((new_border, new_border_count)),
true,
)
@@ -444,7 +316,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
// ...||...
// ... |...
right_border_count.left = left_border_count.left;
(Some(left_border.clone()), None, true)
(Some(*left_border), None, true)
} else if right_border_count.left + left_border_count.left
<= self.current_bucket_size()
&& far_left_neighbor.is_some()
@@ -455,7 +327,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
fln_count.right += right_border_count.left;
right_border_count.left = fln_count.right;
}
(Some(left_border.clone()), None, true)
(Some(*left_border), None, true)
} else {
// ...|..|...
// ... |.|...
@@ -468,7 +340,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
},
);
(
Some(left_border.clone()),
Some(*left_border),
Some((new_border, new_border_count)),
true,
)
@@ -480,7 +352,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
// ...||...
// ...| ...
left_border_count.right = right_border_count.left;
(Some(right_border.clone()), None, true)
(Some(*right_border), None, true)
} else if left_border_count.right + right_border_count.right
<= self.current_bucket_size()
&& far_right_neighbor.is_some()
@@ -491,7 +363,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
frn_count.left += left_border_count.right;
left_border_count.right = frn_count.left;
}
(Some(right_border.clone()), None, true)
(Some(*right_border), None, true)
} else {
// ...|..|...
// ...|.| ...
@@ -504,7 +376,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
},
);
(
Some(right_border.clone()),
Some(*right_border),
Some((new_border, new_border_count)),
true,
)
@@ -570,12 +442,12 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
let (mut close_neighbors, (mut far_left_neighbor, mut far_right_neighbor)) = {
let mut left_iterator = self
.borders
.range((Unbounded, Included(val.clone())))
.map(|(k, v)| (k.clone(), v.clone()));
.range((Unbounded, Included(val)))
.map(|(k, v)| (*k, v.clone()));
let mut right_iterator = self
.borders
.range((Excluded(val.clone()), Unbounded))
.map(|(k, v)| (k.clone(), v.clone()));
.range((Excluded(val), Unbounded))
.map(|(k, v)| (*k, v.clone()));
(
(left_iterator.next_back(), right_iterator.next()),
(left_iterator.next_back(), right_iterator.next()),
@@ -606,10 +478,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
if let Some((_frn, frn_count)) = &mut far_right_neighbor {
frn_count.left = new_count;
}
(
Some(right_border.clone()),
Some((new_border, new_border_count)),
)
(Some(*right_border), Some((new_border, new_border_count)))
}
}
(Some((left_border, left_border_count)), None) => {
@@ -634,10 +503,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
if let Some((_fln, fln_count)) = &mut far_left_neighbor {
fln_count.right = new_count
}
(
Some(left_border.clone()),
Some((new_border, new_border_count)),
)
(Some(*left_border), Some((new_border, new_border_count)))
}
}
(Some((left_border, left_border_count)), Some((right_border, right_border_count))) => {
@@ -669,10 +535,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
if let Some((_fln, fln_count)) = &mut far_left_neighbor {
fln_count.right = new_border_count.left
}
(
Some(left_border.clone()),
Some((new_border, new_border_count)),
)
(Some(*left_border), Some((new_border, new_border_count)))
} else {
// Can't be moved anymore, create an additional one
// ...|..x.........|...
@@ -701,10 +564,7 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
if let Some((_frn, frn_count)) = &mut far_right_neighbor {
frn_count.left = new_border_count.right
}
(
Some(right_border.clone()),
Some((new_border, new_border_count)),
)
(Some(*right_border), Some((new_border, new_border_count)))
} else {
// Can't be moved anymore, create a new one
// 1: ...|........x...|...
lib/segment/src/index/field_index/index_selector.rs
+1 -1
View File
@@ -5,7 +5,6 @@ use gridstore::Blob;
use super::bool_index::BoolIndex;
use super::bool_index::mutable_bool_index::MutableBoolIndex;
use super::geo_index::{GeoMapIndexGridstoreBuilder, GeoMapIndexMmapBuilder};
use super::histogram::Numericable;
use super::map_index::{MapIndex, MapIndexGridstoreBuilder, MapIndexKey, MapIndexMmapBuilder};
use super::numeric_index::{
Encodable, NumericIndexGridstoreBuilder, NumericIndexIntoInnerValue, NumericIndexMmapBuilder,
@@ -19,6 +18,7 @@ use crate::index::field_index::full_text_index::text_index::FullTextIndex;
use crate::index::field_index::geo_index::GeoMapIndex;
use crate::index::field_index::null_index::MutableNullIndex;
use crate::index::field_index::numeric_index::NumericIndex;
use crate::index::field_index::numeric_point::Numericable;
use crate::index::payload_config::{FullPayloadIndexType, PayloadIndexType};
use crate::json_path::JsonPath;
use crate::types::{PayloadFieldSchema, PayloadSchemaParams};
lib/segment/src/index/field_index/mod.rs
+1
View File
@@ -18,6 +18,7 @@ pub mod map_index;
mod memory_reporter;
pub mod null_index;
pub mod numeric_index;
mod numeric_point;
mod stat_tools;
mod stored_point_to_values;
#[cfg(test)]
lib/segment/src/index/field_index/numeric_index/immutable_numeric_index.rs
+7 -6
View File
@@ -11,8 +11,9 @@ use super::mmap_numeric_index::MmapNumericIndex;
use super::mutable_numeric_index::InMemoryNumericIndex;
use crate::common::Flusher;
use crate::common::operation_error::OperationResult;
use crate::index::field_index::histogram::{Histogram, Numericable, Point};
use crate::index::field_index::histogram::Histogram;
use crate::index::field_index::immutable_point_to_values::ImmutablePointToValues;
use crate::index::field_index::numeric_point::{Numericable, Point};
use crate::index::field_index::stored_point_to_values::StoredValue;
use crate::index::payload_config::StorageType;
@@ -149,7 +150,7 @@ impl<T: Encodable + Numericable> Iterator for NumericKeySortedVecIterator<'_, T>
fn next(&mut self) -> Option<Self::Item> {
while self.start_index < self.end_index {
let key = self.set.data[self.start_index].clone();
let key = self.set.data[self.start_index];
let deleted = self.set.deleted.get_bit(self.start_index).unwrap_or(true);
self.start_index += 1;
if deleted {
@@ -164,7 +165,7 @@ impl<T: Encodable + Numericable> Iterator for NumericKeySortedVecIterator<'_, T>
impl<T: Encodable + Numericable> DoubleEndedIterator for NumericKeySortedVecIterator<'_, T> {
fn next_back(&mut self) -> Option<Self::Item> {
while self.start_index < self.end_index {
let key = self.set.data[self.end_index - 1].clone();
let key = self.set.data[self.end_index - 1];
let deleted = self.set.deleted.get_bit(self.end_index - 1).unwrap_or(true);
self.end_index -= 1;
if deleted {
@@ -353,7 +354,7 @@ where
map: &NumericKeySortedVec<T>,
point: &Point<T>,
) -> Option<Point<T>> {
map.values_range(Bound::Unbounded, Bound::Excluded(point.clone()))
map.values_range(Bound::Unbounded, Bound::Excluded(*point))
.next_back()
}
@@ -361,7 +362,7 @@ where
map: &NumericKeySortedVec<T>,
point: &Point<T>,
) -> Option<Point<T>> {
map.values_range(Bound::Excluded(point.clone()), Bound::Unbounded)
map.values_range(Bound::Excluded(*point), Bound::Unbounded)
.next()
}
@@ -388,7 +389,7 @@ mod tests {
end_bound: Bound<Point<FloatPayloadType>>,
) {
let set1 = key_set
.values_range(start_bound.clone(), end_bound.clone())
.values_range(start_bound, end_bound)
.collect::<Vec<_>>();
let set2 = encoded_map
lib/segment/src/index/field_index/numeric_index/mmap_numeric_index.rs
+6 -5
View File
@@ -20,7 +20,8 @@ use crate::common::Flusher;
use crate::common::buffered_update_bitslice::BufferedUpdateBitSlice;
use crate::common::operation_error::{OperationError, OperationResult};
use crate::common::stored_bitslice::MmapBitSlice;
use crate::index::field_index::histogram::{Histogram, Numericable, Point};
use crate::index::field_index::histogram::Histogram;
use crate::index::field_index::numeric_point::{Numericable, Point};
use crate::index::field_index::stored_point_to_values::{StoredPointToValues, StoredValue};
const PAIRS_PATH: &str = "data.bin";
@@ -63,7 +64,7 @@ impl<T: Encodable + Numericable> Iterator for NumericIndexPairsIterator<'_, T> {
fn next(&mut self) -> Option<Self::Item> {
while self.start_index < self.end_index {
let key = self.pairs[self.start_index].clone();
let key = self.pairs[self.start_index];
let deleted = self.deleted.get(key.idx as usize).unwrap_or(true);
self.start_index += 1;
if deleted {
@@ -78,7 +79,7 @@ impl<T: Encodable + Numericable> Iterator for NumericIndexPairsIterator<'_, T> {
impl<T: Encodable + Numericable> DoubleEndedIterator for NumericIndexPairsIterator<'_, T> {
fn next_back(&mut self) -> Option<Self::Item> {
while self.start_index < self.end_index {
let key = self.pairs[self.end_index - 1].clone();
let key = self.pairs[self.end_index - 1];
let deleted = self.deleted.get(key.idx as usize).unwrap_or(true);
self.end_index -= 1;
if deleted {
@@ -128,7 +129,7 @@ impl<T: Encodable + Numericable + Default + StoredValue> MmapNumericIndex<T> {
let pairs_mmap = unsafe { MmapMut::map_mut(&pairs_file)? };
let mut pairs = unsafe { MmapSlice::<Point<T>>::try_from(pairs_mmap)? };
for (src, dst) in in_memory_index.map.iter().zip(pairs.iter_mut()) {
*dst = src.clone();
*dst = *src;
}
}
@@ -298,7 +299,7 @@ impl<T: Encodable + Numericable + Default + StoredValue> MmapNumericIndex<T> {
end_bound: Bound<Point<T>>,
) -> impl DoubleEndedIterator<Item = (T, PointOffsetType)> + '_ {
self.values_range_iterator(start_bound, end_bound)
.map(|Point { val, idx }| (val, idx))
.map(|Point { val, idx, .. }| (val, idx))
}
pub fn remove_point(&mut self, idx: PointOffsetType) {
lib/segment/src/index/field_index/numeric_index/mod.rs
+2 -2
View File
@@ -27,12 +27,12 @@ use uuid::Uuid;
use self::immutable_numeric_index::ImmutableNumericIndex;
use super::FieldIndexBuilderTrait;
use super::histogram::Point;
use super::stored_point_to_values::StoredValue;
use super::utils::{check_boundaries, value_to_integer};
use crate::common::Flusher;
use crate::common::operation_error::{OperationError, OperationResult};
use crate::index::field_index::histogram::{Histogram, Numericable};
use crate::index::field_index::histogram::Histogram;
use crate::index::field_index::numeric_point::{Numericable, Point};
use crate::index::field_index::stat_tools::estimate_multi_value_selection_cardinality;
use crate::index::field_index::{
CardinalityEstimation, PayloadBlockCondition, PayloadFieldIndex, PrimaryCondition, ValueIndexer,
lib/segment/src/index/field_index/numeric_index/mutable_numeric_index.rs
+9 -8
View File
@@ -13,7 +13,8 @@ use super::mmap_numeric_index::MmapNumericIndex;
use super::{Encodable, HISTOGRAM_MAX_BUCKET_SIZE, HISTOGRAM_PRECISION};
use crate::common::Flusher;
use crate::common::operation_error::{OperationError, OperationResult};
use crate::index::field_index::histogram::{Histogram, Numericable, Point};
use crate::index::field_index::histogram::Histogram;
use crate::index::field_index::numeric_point::{Numericable, Point};
use crate::index::field_index::stored_point_to_values::StoredValue;
use crate::index::payload_config::StorageType;
@@ -186,15 +187,15 @@ impl<T: Encodable + Numericable + Default> InMemoryNumericIndex<T> {
}
fn add_to_map(map: &mut BTreeSet<Point<T>>, histogram: &mut Histogram<T>, key: Point<T>) {
let was_added = map.insert(key.clone());
let was_added = map.insert(key);
// Histogram works with unique values (idx + value) only, so we need to
// make sure that we don't add the same value twice.
// key is a combination of value + idx, so we can use it to ensure than the pair is unique
if was_added {
histogram.insert(
key,
|x| Self::get_histogram_left_neighbor(map, x.clone()),
|x| Self::get_histogram_right_neighbor(map, x.clone()),
|x| Self::get_histogram_left_neighbor(map, *x),
|x| Self::get_histogram_right_neighbor(map, *x),
);
}
}
@@ -204,18 +205,18 @@ impl<T: Encodable + Numericable + Default> InMemoryNumericIndex<T> {
if was_removed {
histogram.remove(
&key,
|x| Self::get_histogram_left_neighbor(map, x.clone()),
|x| Self::get_histogram_right_neighbor(map, x.clone()),
|x| Self::get_histogram_left_neighbor(map, *x),
|x| Self::get_histogram_right_neighbor(map, *x),
);
}
}
fn get_histogram_left_neighbor(map: &BTreeSet<Point<T>>, key: Point<T>) -> Option<Point<T>> {
map.range((Unbounded, Excluded(key))).next_back().cloned()
map.range((Unbounded, Excluded(key))).next_back().copied()
}
fn get_histogram_right_neighbor(map: &BTreeSet<Point<T>>, key: Point<T>) -> Option<Point<T>> {
map.range((Excluded(key), Unbounded)).next().cloned()
map.range((Excluded(key), Unbounded)).next().copied()
}
pub fn get_histogram(&self) -> &Histogram<T> {
lib/segment/src/index/field_index/numeric_point.rs
+174
View File
@@ -0,0 +1,174 @@
use std::fmt::Debug;
use common::types::PointOffsetType;
use num_traits::Num;
use serde::Serialize;
pub use self::point::Point;
// bytemuck macros expand to code that triggers this clippy lint
// The only reason this is its own module is so that we scope the lint suppression
#[expect(clippy::multiple_bound_locations)]
mod point {
use common::types::PointOffsetType;
use super::Numericable;
#[expect(clippy::derive_ord_xor_partial_ord)]
#[derive(
PartialEq,
PartialOrd,
Debug,
Clone,
Copy,
serde::Serialize,
serde::Deserialize,
bytemuck::Pod,
bytemuck::Zeroable,
)]
#[repr(C, packed)]
pub struct Point<T: Numericable> {
pub val: T,
pub idx: PointOffsetType,
#[serde(skip)]
_padding: T::PointPadding,
}
impl<T: Numericable> Point<T> {
pub fn new(val: T, idx: PointOffsetType) -> Self {
Self {
val,
idx,
_padding: bytemuck::Zeroable::zeroed(),
}
}
}
impl<T: PartialEq + Numericable> Eq for Point<T> {}
impl<T: PartialOrd + Copy + Numericable> Ord for Point<T> {
fn cmp(&self, other: &Point<T>) -> std::cmp::Ordering {
(self.val, self.idx)
.partial_cmp(&(other.val, other.idx))
.unwrap()
}
}
}
/// Calculate the exact padding so that the [`Point<T>`] type would have explicit alignment, so that
/// it is able to derive [`bytemuck::Pod`] and use #[repr(packed)] safely
const fn derive_point_padding<T: bytemuck::Pod>() -> usize {
struct Point<T> {
_t: T,
_idx: PointOffsetType,
}
let align = std::mem::align_of::<Point<T>>();
// Since we are adding padding at the end, we need to ensure that the align of T is larger the one
// of PointOffsetType, so that there is no inter-field padding
assert!(std::mem::align_of::<T>() >= std::mem::align_of::<PointOffsetType>());
align - (std::mem::size_of::<T>() + std::mem::size_of::<PointOffsetType>()) % align
}
/// A trait that should represent common properties of integer and floating point types.
/// In particular, i64 and f64.
pub trait Numericable: Num + PartialEq + PartialOrd + Copy + bytemuck::Pod {
/// This is to be able to derive [`bytemuck::Pod`] for [`Point<T>`], which is required for safe de/serialization.
///
/// Since we need ['Point<T>`] to be repr(packed), this padding must be picked to be the next multiple of the largest
/// field in the struct, which fits the entire struct.
type PointPadding: bytemuck::Pod
+ Debug
+ Default
+ PartialEq
+ PartialOrd
+ Serialize
+ for<'de> serde::Deserialize<'de>;
fn min_value() -> Self;
fn max_value() -> Self;
fn to_f64(self) -> f64;
fn from_f64(x: f64) -> Self;
fn from_u128(x: u128) -> Self;
fn min(self, b: Self) -> Self {
if self < b { self } else { b }
}
fn max(self, b: Self) -> Self {
if self > b { self } else { b }
}
fn abs_diff(self, b: Self) -> Self {
if self > b { self - b } else { b - self }
}
}
impl Numericable for i64 {
type PointPadding = [u8; derive_point_padding::<Self>()];
fn min_value() -> Self {
i64::MIN
}
fn max_value() -> Self {
i64::MAX
}
fn to_f64(self) -> f64 {
self as f64
}
fn from_f64(x: f64) -> Self {
x as Self
}
fn from_u128(x: u128) -> Self {
x as i64
}
fn abs_diff(self, b: Self) -> Self {
i64::abs_diff(self, b) as i64
}
}
impl Numericable for f64 {
type PointPadding = [u8; derive_point_padding::<Self>()];
fn min_value() -> Self {
f64::MIN
}
fn max_value() -> Self {
f64::MAX
}
fn to_f64(self) -> f64 {
self
}
fn from_f64(x: f64) -> Self {
x
}
fn from_u128(x: u128) -> Self {
x as Self
}
}
impl Numericable for u128 {
type PointPadding = [u8; derive_point_padding::<Self>()];
fn min_value() -> Self {
u128::MIN
}
fn max_value() -> Self {
u128::MAX
}
fn to_f64(self) -> f64 {
self as f64
}
fn from_f64(x: f64) -> Self {
x as u128
}
fn from_u128(x: u128) -> Self {
x
}
fn abs_diff(self, b: Self) -> Self {
u128::abs_diff(self, b)
}
}
lib/segment/src/index/field_index/tests/histogram_i64_tests.rs
+3 -5
View File
@@ -4,7 +4,8 @@ use std::collections::Bound::Included;
use common::types::PointOffsetType;
use rand::prelude::SliceRandom;
use crate::index::field_index::histogram::{Histogram, Point};
use crate::index::field_index::histogram::Histogram;
use crate::index::field_index::numeric_point::Point;
use crate::index::field_index::tests::histogram_test_utils::print_results;
use crate::index::field_index::tests::histogram_tests::{build_histogram, count_range};
@@ -191,10 +192,7 @@ fn test_build_i64_histogram() {
// let points = (0..100000).map(|i| Point { val: rnd.random_range(-10.0..10.0), idx: i }).collect_vec();
let points: Vec<_> = (0..num_samples)
.map(|i| Point {
val: rand::random::<i64>(),
idx: i,
})
.map(|i| Point::new(rand::random::<i64>(), i))
.collect();
let (histogram, points_index) = build_histogram(max_bucket_size, precision, points);
lib/segment/src/index/field_index/tests/histogram_test_utils.rs
+2 -1
View File
@@ -6,7 +6,8 @@ use std::io::Write;
use serde::Serialize;
use serde::de::DeserializeOwned;
use crate::index::field_index::histogram::{Histogram, Numericable, Point};
use crate::index::field_index::histogram::Histogram;
use crate::index::field_index::numeric_point::{Numericable, Point};
pub fn print_results<T: Numericable + Serialize + DeserializeOwned + Display>(
points_index: &BTreeSet<Point<T>>,
lib/segment/src/index/field_index/tests/histogram_tests.rs
+24 -22
View File
@@ -9,13 +9,21 @@ use rand_distr::StandardNormal;
use serde::Serialize;
use serde::de::DeserializeOwned;
use crate::index::field_index::histogram::{Histogram, Numericable, Point};
use crate::index::field_index::histogram::Histogram;
use crate::index::field_index::numeric_point::{Numericable, Point};
use crate::index::field_index::tests::histogram_test_utils::print_results;
pub fn count_range<T: PartialOrd>(points_index: &BTreeSet<Point<T>>, a: T, b: T) -> usize {
pub fn count_range<T: PartialOrd + Numericable>(
points_index: &BTreeSet<Point<T>>,
a: T,
b: T,
) -> usize {
points_index
.iter()
.filter(|x| a <= x.val && x.val <= b)
.filter(|x| {
let v = x.val;
a <= v && v <= b
})
.count()
}
@@ -28,9 +36,11 @@ fn test_build_histogram_small() {
// let points = (0..100000).map(|i| Point { val: rnd.random_range(-10.0..10.0), idx: i }).collect_vec();
let points = (0..num_samples)
.map(|i| Point {
val: f64::round(rnd.sample::<f64, _>(StandardNormal) * 10.0),
idx: i % num_samples / 2,
.map(|i| {
Point::new(
f64::round(rnd.sample::<f64, _>(StandardNormal) * 10.0),
i % num_samples / 2,
)
})
.collect_vec();
@@ -39,10 +49,10 @@ fn test_build_histogram_small() {
let mut histogram = Histogram::new(max_bucket_size, precision);
for point in &points {
points_index.insert(point.clone());
points_index.insert(*point);
// print_results(&points_index, &histogram, Some(point.clone()));
histogram.insert(
point.clone(),
*point,
|x| {
points_index
.range((Unbounded, Excluded(x)))
@@ -54,7 +64,7 @@ fn test_build_histogram_small() {
}
for point in &points {
print_results(&points_index, &histogram, Some(point.clone()));
print_results(&points_index, &histogram, Some(*point));
points_index.remove(point);
histogram.remove(
point,
@@ -193,7 +203,7 @@ pub fn build_histogram<T: Numericable + Serialize + DeserializeOwned + std::fmt:
let read_counter = Cell::new(0);
for point in points {
points_index.insert(point.clone());
points_index.insert(point);
// print_results(&points_index, &histogram, Some(point.clone()));
histogram.insert(
point,
@@ -226,10 +236,8 @@ fn test_build_histogram_round() {
let mut rnd = StdRng::seed_from_u64(42);
// let points = (0..100000).map(|i| Point { val: rnd.random_range(-10.0..10.0), idx: i }).collect_vec();
let points = (0..num_samples).map(|i| Point {
val: f64::round(rnd.sample::<f64, _>(StandardNormal) * 100.0),
idx: i,
});
let points = (0..num_samples)
.map(|i| Point::new(f64::round(rnd.sample::<f64, _>(StandardNormal) * 100.0), i));
let (histogram, points_index) = build_histogram(max_bucket_size, precision, points.collect());
request_histogram(&histogram, &points_index);
@@ -244,10 +252,7 @@ fn test_build_histogram() {
// let points = (0..100000).map(|i| Point { val: rnd.random_range(-10.0..10.0), idx: i }).collect_vec();
let points = (0..num_samples)
.map(|i| Point {
val: rnd.sample(StandardNormal),
idx: i,
})
.map(|i| Point::new(rnd.sample(StandardNormal), i))
.collect_vec();
let (histogram, points_index) = build_histogram(max_bucket_size, precision, points);
@@ -263,10 +268,7 @@ fn test_save_load_histogram() {
let mut rnd = StdRng::seed_from_u64(42);
let points = (0..num_samples)
.map(|i| Point {
val: rnd.random_range(-10.0..10.0),
idx: i,
})
.map(|i| Point::new(rnd.random_range(-10.0..10.0), i))
.collect_vec();
let (histogram, _) = build_histogram(max_bucket_size, precision, points);