use std::borrow::Cow;
use std::cmp::Ordering;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::fmt::{self, Display, Formatter};
use std::hash::{self, Hash, Hasher};
use std::mem;
use std::ops::Deref;
use std::rc::Rc;
use std::str::FromStr;
use std::sync::Arc;
use ahash::AHashSet;
use bytemuck::{Pod, Zeroable};
use crate::common::stable_hash::StableHash;
use crate::common::types::{PointOffsetType, ScoreType};
use ecow::EcoString;
use fnv::FnvBuildHasher;
use geo::{Contains, Coord, Distance as GeoDistance, Haversine, LineString, Point, Polygon};
use indexmap::IndexSet;
use itertools::Itertools;
use num_derive::FromPrimitive;
use ordered_float::OrderedFloat;
use schemars::JsonSchema;
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value};
use strum::{EnumIter, EnumString};
use uuid::Uuid;
use validator::{Validate, ValidationError, ValidationErrors};
use zerocopy::native_endian::U64;
use crate::segment::common::anonymize::Anonymize;
use crate::segment::common::operation_error::{OperationError, OperationResult};
use crate::segment::common::utils::{self, MaybeOneOrMany, MultiValue};
use crate::segment::data_types::index::{
BoolIndexParams, DatetimeIndexParams, FloatIndexParams, GeoIndexParams, IntegerIndexParams,
KeywordIndexParams, TextIndexParams, UuidIndexParams,
};
use crate::segment::data_types::modifier::Modifier;
use crate::segment::data_types::order_by::OrderValue;
use crate::segment::data_types::primitive::PrimitiveVectorElement;
use crate::segment::data_types::tiny_map::TinyMap;
use crate::segment::data_types::vectors::{DenseVector, VectorStructInternal};
use crate::segment::index::field_index::CardinalityEstimation;
use crate::segment::index::sparse_index::sparse_index_config::SparseIndexConfig;
use crate::segment::json_path::JsonPath;
use crate::segment::spaces::metric::{Metric, MetricPostProcessing};
use crate::segment::spaces::simple::{CosineMetric, DotProductMetric, EuclidMetric, ManhattanMetric};
use crate::segment::types::utils::unordered_hash_unique;
use crate::segment::utils::maybe_arc::MaybeArc;
pub type PayloadKeyType = JsonPath;
pub type PayloadKeyTypeRef<'a> = &'a JsonPath;
pub type SeqNumberType = u64;
pub type FloatPayloadType = f64;
pub type IntPayloadType = i64;
pub type DateTimePayloadType = DateTimeWrapper;
pub type UuidPayloadType = Uuid;
pub type UuidIntType = u128;
pub type VectorName = str;
pub type VectorNameBuf = String;
#[derive(Clone, Copy, Serialize, JsonSchema, Debug, PartialEq, Eq, PartialOrd, Hash)]
#[serde(transparent)]
pub struct DateTimeWrapper(pub chrono::DateTime<chrono::Utc>);
impl DateTimeWrapper {
pub fn timestamp(&self) -> i64 {
self.0.timestamp_micros()
}
pub fn from_timestamp(ts: i64) -> Option<Self> {
Some(Self(chrono::DateTime::from_timestamp_micros(ts)?))
}
}
impl<'de> Deserialize<'de> for DateTimePayloadType {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let str_datetime: Cow<'de, str> = Cow::deserialize(deserializer)?;
match DateTimePayloadType::from_str(str_datetime.as_ref()) {
Ok(datetime) => Ok(datetime),
Err(_) => Err(serde::de::Error::custom(format!(
"'{str_datetime}' does not match accepted datetime format (RFC3339). Example: 2014-01-01T00:00:00Z"
))),
}
}
}
impl FromStr for DateTimePayloadType {
type Err = chrono::ParseError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
if let Ok(datetime) = chrono::DateTime::parse_from_rfc3339(s)
.or_else(|_| chrono::DateTime::from_str(s))
.or_else(|_| chrono::DateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.f%#z"))
.or_else(|_| chrono::DateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f%#z"))
.map(|dt| chrono::DateTime::<chrono::Utc>::from(dt).into())
{
return Ok(datetime);
}
let datetime = chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.f")
.or_else(|_| chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f"))
.or_else(|_| chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M"))
.or_else(|_| chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M"))
.or_else(|_| chrono::NaiveDate::parse_from_str(s, "%Y-%m-%d").map(Into::into))?;
let datetime_utc = datetime.and_utc().into();
Ok(datetime_utc)
}
}
impl Display for DateTimePayloadType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
impl From<chrono::DateTime<chrono::Utc>> for DateTimePayloadType {
fn from(dt: chrono::DateTime<chrono::Utc>) -> Self {
DateTimeWrapper(dt)
}
}
fn id_num_example() -> u64 {
42
}
fn id_uuid_example() -> String {
"550e8400-e29b-41d4-a716-446655440000".to_string()
}
#[derive(Debug, Serialize, Copy, Clone, PartialEq, Eq, Hash, Ord, PartialOrd, JsonSchema)]
#[serde(untagged)]
pub enum ExtendedPointId {
#[schemars(example = "id_num_example")]
NumId(u64),
#[schemars(example = "id_uuid_example")]
Uuid(Uuid),
}
impl StableHash for ExtendedPointId {
fn stable_hash<W: FnMut(&[u8])>(&self, write: &mut W) {
match self {
ExtendedPointId::NumId(num) => {
0u64.stable_hash(write); num.stable_hash(write);
}
ExtendedPointId::Uuid(uuid) => {
1u64.stable_hash(write);
uuid.as_bytes().len().stable_hash(write); write(uuid.as_bytes());
}
}
}
}
impl ExtendedPointId {
#[cfg(any(test, feature = "testing"))]
pub fn as_u64(&self) -> u64 {
match self {
ExtendedPointId::NumId(num) => *num,
ExtendedPointId::Uuid(_) => panic!("Cannot convert UUID to u64"),
}
}
pub fn is_uuid(&self) -> bool {
matches!(self, ExtendedPointId::Uuid(..))
}
}
impl std::fmt::Display for ExtendedPointId {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
ExtendedPointId::NumId(idx) => write!(f, "{idx}"),
ExtendedPointId::Uuid(uuid) => write!(f, "{uuid}"),
}
}
}
impl From<u64> for ExtendedPointId {
fn from(idx: u64) -> Self {
ExtendedPointId::NumId(idx)
}
}
impl FromStr for ExtendedPointId {
type Err = ();
fn from_str(s: &str) -> Result<Self, Self::Err> {
let try_num: Result<u64, _> = s.parse();
if let Ok(num) = try_num {
return Ok(Self::NumId(num));
}
let try_uuid = Uuid::from_str(s);
if let Ok(uuid) = try_uuid {
return Ok(Self::Uuid(uuid));
}
Err(())
}
}
impl<'de> serde::Deserialize<'de> for ExtendedPointId {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let value = serde_value::Value::deserialize(deserializer)?;
if let Ok(num) = value.clone().deserialize_into() {
return Ok(ExtendedPointId::NumId(num));
}
if let Ok(uuid) = value.clone().deserialize_into() {
return Ok(ExtendedPointId::Uuid(uuid));
}
let value = crate::segment::utils::fmt::SerdeValue(&value);
Err(serde::de::Error::custom(format!(
"value {value} is not a valid point ID, \
valid values are either an unsigned integer or a UUID",
)))
}
}
pub type PointIdType = ExtendedPointId;
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Ord, PartialOrd)]
pub enum CompactExtendedPointId {
NumId(U64),
Uuid(Uuid),
}
impl From<ExtendedPointId> for CompactExtendedPointId {
fn from(id: ExtendedPointId) -> Self {
match id {
ExtendedPointId::NumId(num) => CompactExtendedPointId::NumId(U64::new(num)),
ExtendedPointId::Uuid(uuid) => CompactExtendedPointId::Uuid(uuid),
}
}
}
impl From<CompactExtendedPointId> for ExtendedPointId {
fn from(id: CompactExtendedPointId) -> Self {
match id {
CompactExtendedPointId::NumId(num) => ExtendedPointId::NumId(num.get()),
CompactExtendedPointId::Uuid(uuid) => ExtendedPointId::Uuid(uuid),
}
}
}
#[derive(
Debug,
Deserialize,
Serialize,
JsonSchema,
Clone,
Copy,
FromPrimitive,
PartialEq,
Eq,
Hash,
EnumString,
EnumIter,
)]
pub enum Distance {
Cosine,
Euclid,
Dot,
Manhattan,
}
impl Distance {
pub fn postprocess_score(&self, score: ScoreType) -> ScoreType {
match self {
Distance::Cosine => CosineMetric::postprocess(score),
Distance::Euclid => EuclidMetric::postprocess(score),
Distance::Dot => DotProductMetric::postprocess(score),
Distance::Manhattan => ManhattanMetric::postprocess(score),
}
}
pub fn preprocess_vector<T: PrimitiveVectorElement>(&self, vector: DenseVector) -> DenseVector
where
CosineMetric: Metric<T>,
EuclidMetric: Metric<T>,
DotProductMetric: Metric<T>,
ManhattanMetric: Metric<T>,
{
match self {
Distance::Cosine => CosineMetric::preprocess(vector),
Distance::Euclid => EuclidMetric::preprocess(vector),
Distance::Dot => DotProductMetric::preprocess(vector),
Distance::Manhattan => ManhattanMetric::preprocess(vector),
}
}
pub fn distance_order(&self) -> Order {
match self {
Distance::Cosine | Distance::Dot => Order::LargeBetter,
Distance::Euclid | Distance::Manhattan => Order::SmallBetter,
}
}
pub fn is_ordered(&self, left: ScoreType, right: ScoreType) -> bool {
match self.distance_order() {
Order::LargeBetter => left >= right,
Order::SmallBetter => left <= right,
}
}
pub fn check_threshold(&self, score: ScoreType, threshold: ScoreType) -> bool {
match self.distance_order() {
Order::LargeBetter => score > threshold,
Order::SmallBetter => score < threshold,
}
}
}
#[derive(Debug, PartialEq, Clone, Copy)]
pub enum Order {
LargeBetter,
SmallBetter,
}
#[derive(Clone, Debug)]
pub struct ScoredPoint {
pub id: PointIdType,
pub version: SeqNumberType,
pub score: ScoreType,
pub payload: Option<Payload>,
pub vector: Option<VectorStructInternal>,
pub shard_key: Option<ShardKey>,
pub order_value: Option<OrderValue>,
}
impl Eq for ScoredPoint {}
impl Ord for ScoredPoint {
fn cmp(&self, other: &Self) -> Ordering {
match (&self.order_value, &other.order_value) {
(None, None) => OrderedFloat(self.score).cmp(&OrderedFloat(other.score)),
(Some(_), None) => Ordering::Greater,
(None, Some(_)) => Ordering::Less,
(Some(self_order), Some(other_order)) => self_order.cmp(other_order),
}
}
}
impl PartialOrd for ScoredPoint {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl PartialEq for ScoredPoint {
fn eq(&self, other: &Self) -> bool {
(self.id, &self.score) == (other.id, &other.score)
}
}
#[derive(Debug, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum SegmentType {
Plain,
Indexed,
Special,
}
#[derive(Debug, Serialize, JsonSchema, Clone, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub struct PayloadIndexInfo {
pub data_type: PayloadSchemaType,
#[serde(skip_serializing_if = "Option::is_none")]
pub params: Option<PayloadSchemaParams>,
pub points: usize,
}
impl PayloadIndexInfo {
pub fn new(field_type: PayloadFieldSchema, points_count: usize) -> Self {
match field_type {
PayloadFieldSchema::FieldType(data_type) => PayloadIndexInfo {
data_type,
params: None,
points: points_count,
},
PayloadFieldSchema::FieldParams(schema_params) => PayloadIndexInfo {
data_type: schema_params.kind(),
params: Some(schema_params),
points: points_count,
},
}
}
}
#[derive(Debug, Serialize, JsonSchema, Clone, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub struct VectorDataInfo {
pub num_vectors: usize,
pub num_indexed_vectors: usize,
pub num_deleted_vectors: usize,
}
#[derive(Debug, Serialize, JsonSchema, Clone, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub struct SegmentInfo {
pub uuid: Uuid,
pub segment_type: SegmentType,
pub num_vectors: usize,
pub num_points: usize,
pub num_deferred_points: Option<usize>,
pub num_deleted_deferred_points: Option<usize>,
pub num_indexed_vectors: usize,
pub num_deleted_vectors: usize,
pub vectors_size_bytes: usize,
pub payloads_size_bytes: usize,
pub ram_usage_bytes: usize,
pub disk_usage_bytes: usize,
pub is_appendable: bool,
pub index_schema: HashMap<PayloadKeyType, PayloadIndexInfo>,
pub vector_data: HashMap<String, VectorDataInfo>,
#[serde(skip_serializing_if = "Option::is_none")]
pub deferred_internal_id: Option<PointOffsetType>,
}
#[derive(Debug, Default)]
pub struct SizeStats {
pub num_vectors: usize,
pub num_vectors_by_name: TinyMap<VectorNameBuf, usize>,
pub vectors_size_bytes: usize,
pub payloads_size_bytes: usize,
pub num_points: usize,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Validate, Clone, Copy, PartialEq, Default)]
#[serde(rename_all = "snake_case")]
pub struct QuantizationSearchParams {
#[serde(default = "default_quantization_ignore_value")]
pub ignore: bool,
#[serde(default)]
#[serde(skip_serializing_if = "Option::is_none")]
pub rescore: Option<bool>,
#[serde(default = "default_quantization_oversampling_value")]
#[validate(range(min = 1.0))]
#[serde(skip_serializing_if = "Option::is_none")]
pub oversampling: Option<f64>,
}
impl Hash for QuantizationSearchParams {
fn hash<H: Hasher>(&self, state: &mut H) {
let Self {
ignore,
rescore,
oversampling,
} = self;
ignore.hash(state);
rescore.hash(state);
oversampling.map(OrderedFloat).hash(state);
}
}
pub const fn default_quantization_ignore_value() -> bool {
false
}
pub const fn default_quantization_oversampling_value() -> Option<f64> {
None
}
pub const ACORN_MAX_SELECTIVITY_DEFAULT: f64 = 0.4;
#[derive(
Debug, Deserialize, Serialize, JsonSchema, Validate, Clone, Copy, PartialEq, Default, Hash,
)]
#[serde(rename_all = "snake_case")]
pub struct AcornSearchParams {
#[serde(default)]
pub enable: bool,
#[serde(default)]
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 0.0, max = 1.0))]
pub max_selectivity: Option<OrderedFloat<f64>>,
}
#[derive(
Debug, Deserialize, Serialize, JsonSchema, Validate, Copy, Clone, PartialEq, Default, Hash,
)]
#[serde(rename_all = "snake_case")]
pub struct SearchParams {
#[serde(skip_serializing_if = "Option::is_none")]
pub hnsw_ef: Option<usize>,
#[serde(default)]
pub exact: bool,
#[serde(default)]
#[validate(nested)]
#[serde(skip_serializing_if = "Option::is_none")]
pub quantization: Option<QuantizationSearchParams>,
#[serde(default)]
pub indexed_only: bool,
#[serde(default)]
#[validate(nested)]
#[serde(skip_serializing_if = "Option::is_none")]
pub acorn: Option<AcornSearchParams>,
}
#[derive(Debug, Deserialize, Validate, Clone, PartialEq, Eq)]
pub struct VectorsConfigDefaults {
#[serde(default)]
pub on_disk: Option<bool>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
#[serde(tag = "type", content = "options")]
pub enum Indexes {
Plain {},
Hnsw(HnswConfig),
}
impl Indexes {
pub fn is_indexed(&self) -> bool {
match self {
Indexes::Plain {} => false,
Indexes::Hnsw(_) => true,
}
}
pub fn is_on_disk(&self) -> bool {
match self {
Indexes::Plain {} => false,
Indexes::Hnsw(config) => config.on_disk.unwrap_or_default(),
}
}
}
#[derive(
Copy, Clone, Debug, Eq, PartialEq, Deserialize, Serialize, JsonSchema, Validate,
)]
#[serde(rename_all = "snake_case")]
pub struct HnswConfig {
pub m: usize,
#[validate(range(min = 4))]
pub ef_construct: usize,
#[serde(alias = "full_scan_threshold_kb")]
pub full_scan_threshold: usize,
#[serde(default = "default_max_indexing_threads")]
pub max_indexing_threads: usize,
#[serde(default, skip_serializing_if = "Option::is_none")] pub on_disk: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")] pub payload_m: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub inline_storage: Option<bool>,
}
impl HnswConfig {
pub fn mismatch_requires_rebuild(&self, other: &Self) -> bool {
let HnswConfig {
m,
ef_construct,
full_scan_threshold,
max_indexing_threads: _,
payload_m,
on_disk,
inline_storage,
} = *self;
m != other.m
|| ef_construct != other.ef_construct
|| full_scan_threshold != other.full_scan_threshold
|| payload_m != other.payload_m
|| on_disk != other.on_disk
|| inline_storage != other.inline_storage
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Validate, Clone)]
#[serde(rename_all = "snake_case", default)]
pub struct HnswGlobalConfig {
#[validate(range(min = 0.0, max = 1.0))]
pub healing_threshold: f64,
}
impl Default for HnswGlobalConfig {
fn default() -> Self {
Self {
healing_threshold: 0.3,
}
}
}
const fn default_max_indexing_threads() -> usize {
0
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq, Hash)]
#[serde(rename_all = "lowercase")]
pub enum CompressionRatio {
X4,
X8,
X16,
X32,
X64,
}
#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Hash, Deserialize, Serialize, JsonSchema)]
#[serde(rename_all = "lowercase")]
pub enum ScalarType {
#[default]
Int8,
}
#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, JsonSchema, Validate)]
#[serde(rename_all = "snake_case")]
pub struct ScalarQuantizationConfig {
pub r#type: ScalarType,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 0.5, max = 1.0))]
pub quantile: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub always_ram: Option<bool>,
}
impl ScalarQuantizationConfig {
pub fn mismatch_requires_rebuild(&self, other: &Self) -> bool {
self != other
}
}
#[derive(Clone, Debug, Eq, PartialEq, Hash, Deserialize, Serialize, JsonSchema, Validate)]
pub struct ScalarQuantization {
#[validate(nested)]
pub scalar: ScalarQuantizationConfig,
}
#[derive(Clone, Debug, Eq, PartialEq, Hash, Deserialize, Serialize, JsonSchema, Validate)]
#[serde(rename_all = "snake_case")]
pub struct ProductQuantizationConfig {
pub compression: CompressionRatio,
#[serde(skip_serializing_if = "Option::is_none")]
pub always_ram: Option<bool>,
}
impl ProductQuantizationConfig {
pub fn mismatch_requires_rebuild(&self, other: &Self) -> bool {
self != other
}
}
#[derive(Clone, Debug, Eq, PartialEq, Hash, Deserialize, Serialize, JsonSchema, Validate)]
pub struct ProductQuantization {
#[validate(nested)]
pub product: ProductQuantizationConfig,
}
impl Hash for ScalarQuantizationConfig {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.always_ram.hash(state);
self.r#type.hash(state);
}
}
impl Eq for ScalarQuantizationConfig {}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq, Hash, Default)]
#[serde(rename_all = "snake_case")]
pub enum BinaryQuantizationEncoding {
#[default]
OneBit,
TwoBits,
OneAndHalfBits,
}
#[derive(Clone, Debug, Eq, PartialEq, Hash, Deserialize, Serialize, JsonSchema, Validate)]
#[serde(rename_all = "snake_case")]
pub struct BinaryQuantizationConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub always_ram: Option<bool>,
#[serde(default)]
#[serde(skip_serializing_if = "Option::is_none")]
pub encoding: Option<BinaryQuantizationEncoding>,
#[serde(default)]
#[serde(skip_serializing_if = "Option::is_none")]
pub query_encoding: Option<BinaryQuantizationQueryEncoding>,
}
#[derive(Clone, Debug, Eq, PartialEq, Hash, Deserialize, Serialize, JsonSchema, Validate)]
pub struct BinaryQuantization {
#[validate(nested)]
pub binary: BinaryQuantizationConfig,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq, Hash, Default)]
#[serde(rename_all = "snake_case")]
pub enum TurboQuantBitSize {
Bits1,
Bits1_5,
Bits2,
#[default]
Bits4,
}
#[derive(Clone, Debug, Eq, PartialEq, Hash, Deserialize, Serialize, JsonSchema, Validate)]
#[serde(rename_all = "snake_case")]
pub struct TurboQuantQuantizationConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub always_ram: Option<bool>,
#[serde(default)]
#[serde(skip_serializing_if = "Option::is_none")]
pub bits: Option<TurboQuantBitSize>,
}
#[derive(Clone, Debug, Eq, PartialEq, Hash, Deserialize, Serialize, JsonSchema, Validate)]
pub struct TurboQuantization {
#[validate(nested)]
pub turbo: TurboQuantQuantizationConfig,
}
#[derive(Clone, Debug, Eq, PartialEq, Hash, Deserialize, Serialize, JsonSchema, )]
#[serde(untagged, rename_all = "snake_case")]
pub enum QuantizationConfig {
Scalar(ScalarQuantization),
Product(ProductQuantization),
Binary(BinaryQuantization),
Turbo(TurboQuantization),
}
impl QuantizationConfig {
pub fn for_appendable_segment(opt: Option<&Self>) -> Option<Self> {
let appendable = crate::common::flags::feature_flags().appendable_quantization;
opt.filter(|q| appendable && q.supports_appendable())
.cloned()
}
pub fn mismatch_requires_rebuild(&self, other: &Self) -> bool {
self != other
}
pub fn supports_appendable(&self) -> bool {
matches!(
self,
QuantizationConfig::Binary(_) | QuantizationConfig::Turbo(_)
)
}
pub fn always_ram(&self) -> bool {
match self {
QuantizationConfig::Scalar(s) => s.scalar.always_ram == Some(true),
QuantizationConfig::Product(p) => p.product.always_ram == Some(true),
QuantizationConfig::Binary(b) => b.binary.always_ram == Some(true),
QuantizationConfig::Turbo(t) => t.turbo.always_ram == Some(true),
}
}
}
impl Validate for QuantizationConfig {
fn validate(&self) -> Result<(), ValidationErrors> {
match self {
QuantizationConfig::Scalar(scalar) => scalar.validate(),
QuantizationConfig::Product(product) => product.validate(),
QuantizationConfig::Binary(binary) => binary.validate(),
QuantizationConfig::Turbo(turbo) => turbo.validate(),
}
}
}
#[derive(
Default, Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq, Hash,
)]
#[serde(rename_all = "lowercase")]
pub enum BinaryQuantizationQueryEncoding {
#[default]
Default,
Binary,
Scalar4Bits,
Scalar8Bits,
}
impl From<ScalarQuantizationConfig> for QuantizationConfig {
fn from(config: ScalarQuantizationConfig) -> Self {
QuantizationConfig::Scalar(ScalarQuantization { scalar: config })
}
}
impl From<ProductQuantizationConfig> for QuantizationConfig {
fn from(config: ProductQuantizationConfig) -> Self {
QuantizationConfig::Product(ProductQuantization { product: config })
}
}
impl From<BinaryQuantizationConfig> for QuantizationConfig {
fn from(config: BinaryQuantizationConfig) -> Self {
QuantizationConfig::Binary(BinaryQuantization { binary: config })
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Validate, Clone, PartialEq, Default, Hash)]
pub struct StrictModeSparse {
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 1))]
pub max_length: Option<usize>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Validate, Clone, PartialEq, Default, Hash)]
#[schemars(deny_unknown_fields)]
pub struct StrictModeSparseConfig {
#[validate(nested)]
#[serde(flatten)]
pub config: BTreeMap<VectorNameBuf, StrictModeSparse>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Default)]
#[schemars(deny_unknown_fields)]
pub struct StrictModeSparseConfigOutput {
#[serde(flatten)]
pub config: BTreeMap<VectorNameBuf, StrictModeSparseOutput>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Default)]
pub struct StrictModeSparseOutput {
#[serde(skip_serializing_if = "Option::is_none")]
pub max_length: Option<usize>,
}
impl From<StrictModeSparseConfig> for StrictModeSparseConfigOutput {
fn from(config: StrictModeSparseConfig) -> Self {
let StrictModeSparseConfig { config } = config;
let mut new_config = StrictModeSparseConfigOutput::default();
for (key, value) in config {
new_config
.config
.insert(key, StrictModeSparseOutput::from(value));
}
new_config
}
}
impl From<StrictModeSparse> for StrictModeSparseOutput {
fn from(config: StrictModeSparse) -> Self {
let StrictModeSparse { max_length } = config;
StrictModeSparseOutput { max_length }
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Validate, Clone, PartialEq, Default, Hash)]
pub struct StrictModeMultivector {
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 1))]
pub max_vectors: Option<usize>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Validate, Clone, PartialEq, Default, Hash)]
#[schemars(deny_unknown_fields)]
pub struct StrictModeMultivectorConfig {
#[validate(nested)]
#[serde(flatten)]
pub config: BTreeMap<VectorNameBuf, StrictModeMultivector>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Default)]
#[schemars(deny_unknown_fields)]
pub struct StrictModeMultivectorConfigOutput {
#[serde(flatten)]
pub config: BTreeMap<VectorNameBuf, StrictModeMultivectorOutput>,
}
impl From<StrictModeMultivectorConfig> for StrictModeMultivectorConfigOutput {
fn from(config: StrictModeMultivectorConfig) -> Self {
let StrictModeMultivectorConfig { config } = config;
let mut new_config = StrictModeMultivectorConfigOutput::default();
for (key, value) in config {
new_config
.config
.insert(key, StrictModeMultivectorOutput::from(value));
}
new_config
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Default)]
pub struct StrictModeMultivectorOutput {
#[serde(skip_serializing_if = "Option::is_none")]
pub max_vectors: Option<usize>,
}
impl From<StrictModeMultivector> for StrictModeMultivectorOutput {
fn from(config: StrictModeMultivector) -> Self {
let StrictModeMultivector { max_vectors } = config;
StrictModeMultivectorOutput { max_vectors }
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Validate, Clone, PartialEq, Default)]
pub struct StrictModeConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub enabled: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 1))]
pub max_query_limit: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 1))]
pub max_timeout: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub unindexed_filtering_retrieve: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub unindexed_filtering_update: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub search_max_hnsw_ef: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub search_allow_exact: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub search_max_oversampling: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upsert_max_batchsize: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub search_max_batchsize: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_collection_vector_size_bytes: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 1))]
pub read_rate_limit: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 1))]
pub write_rate_limit: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_collection_payload_size_bytes: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 1))]
pub max_points_count: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub filter_max_conditions: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub condition_max_size: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(nested)]
pub multivector_config: Option<StrictModeMultivectorConfig>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(nested)]
pub sparse_config: Option<StrictModeSparseConfig>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 0))]
pub max_payload_index_count: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 1, max = 100))]
pub max_resident_memory_percent: Option<u8>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 1, max = 100))]
pub max_disk_usage_percent: Option<u8>,
}
impl Eq for StrictModeConfig {}
impl Hash for StrictModeConfig {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
let Self {
enabled,
max_query_limit,
max_timeout,
unindexed_filtering_retrieve,
unindexed_filtering_update,
search_max_hnsw_ef,
search_allow_exact,
search_max_oversampling: _,
upsert_max_batchsize,
search_max_batchsize,
max_collection_vector_size_bytes,
read_rate_limit,
write_rate_limit,
max_collection_payload_size_bytes,
max_points_count,
filter_max_conditions,
condition_max_size,
multivector_config,
sparse_config,
max_payload_index_count,
max_resident_memory_percent,
max_disk_usage_percent,
} = self;
enabled.hash(state);
max_query_limit.hash(state);
max_timeout.hash(state);
unindexed_filtering_retrieve.hash(state);
unindexed_filtering_update.hash(state);
search_max_hnsw_ef.hash(state);
search_allow_exact.hash(state);
upsert_max_batchsize.hash(state);
search_max_batchsize.hash(state);
max_collection_vector_size_bytes.hash(state);
read_rate_limit.hash(state);
write_rate_limit.hash(state);
max_collection_payload_size_bytes.hash(state);
max_points_count.hash(state);
filter_max_conditions.hash(state);
condition_max_size.hash(state);
multivector_config.hash(state);
sparse_config.hash(state);
max_payload_index_count.hash(state);
max_resident_memory_percent.hash(state);
max_disk_usage_percent.hash(state);
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Default)]
pub struct StrictModeConfigOutput {
#[serde(skip_serializing_if = "Option::is_none")]
pub enabled: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_query_limit: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_timeout: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub unindexed_filtering_retrieve: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub unindexed_filtering_update: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub search_max_hnsw_ef: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub search_allow_exact: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub search_max_oversampling: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upsert_max_batchsize: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub search_max_batchsize: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_collection_vector_size_bytes: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub read_rate_limit: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub write_rate_limit: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_collection_payload_size_bytes: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_points_count: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub filter_max_conditions: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub condition_max_size: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub multivector_config: Option<StrictModeMultivectorConfigOutput>,
#[serde(skip_serializing_if = "Option::is_none")]
pub sparse_config: Option<StrictModeSparseConfigOutput>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_payload_index_count: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_resident_memory_percent: Option<u8>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_disk_usage_percent: Option<u8>,
}
impl From<StrictModeConfig> for StrictModeConfigOutput {
fn from(config: StrictModeConfig) -> Self {
let StrictModeConfig {
enabled,
max_query_limit,
max_timeout,
unindexed_filtering_retrieve,
unindexed_filtering_update,
search_max_hnsw_ef,
search_allow_exact,
search_max_oversampling,
upsert_max_batchsize,
search_max_batchsize,
max_collection_vector_size_bytes,
read_rate_limit,
write_rate_limit,
max_collection_payload_size_bytes,
max_points_count,
filter_max_conditions,
condition_max_size,
multivector_config,
sparse_config,
max_payload_index_count,
max_resident_memory_percent,
max_disk_usage_percent,
} = config;
Self {
enabled,
max_query_limit,
max_timeout,
unindexed_filtering_retrieve,
unindexed_filtering_update,
search_max_hnsw_ef,
search_allow_exact,
search_max_oversampling,
upsert_max_batchsize,
search_max_batchsize,
max_collection_vector_size_bytes,
read_rate_limit,
write_rate_limit,
max_collection_payload_size_bytes,
max_points_count,
filter_max_conditions,
condition_max_size,
multivector_config: multivector_config.map(StrictModeMultivectorConfigOutput::from),
sparse_config: sparse_config.map(StrictModeSparseConfigOutput::from),
max_payload_index_count,
max_resident_memory_percent,
max_disk_usage_percent,
}
}
}
pub const DEFAULT_HNSW_EF_CONSTRUCT: usize = 100;
impl Default for HnswConfig {
fn default() -> Self {
HnswConfig {
m: 16,
ef_construct: DEFAULT_HNSW_EF_CONSTRUCT,
full_scan_threshold: DEFAULT_FULL_SCAN_THRESHOLD,
max_indexing_threads: 0,
on_disk: Some(false),
payload_m: None,
inline_storage: None,
}
}
}
impl Default for Indexes {
fn default() -> Self {
Indexes::Plain {}
}
}
#[derive( Debug, Deserialize, Serialize, JsonSchema, Copy, Clone, PartialEq, Eq)]
#[serde(tag = "type", content = "options", rename_all = "snake_case")]
pub enum PayloadStorageType {
Mmap,
InRamMmap,
}
#[cfg(any(test, feature = "testing"))]
#[allow(clippy::derivable_impls)]
impl Default for PayloadStorageType {
fn default() -> Self {
PayloadStorageType::Mmap
}
}
impl PayloadStorageType {
pub fn from_on_disk_payload(on_disk: bool) -> Self {
if on_disk { Self::Mmap } else { Self::InRamMmap }
}
pub fn is_on_disk(&self) -> bool {
match self {
PayloadStorageType::Mmap => true,
PayloadStorageType::InRamMmap => false,
}
}
}
#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, JsonSchema, )]
#[serde(rename_all = "snake_case")]
pub struct SegmentConfig {
#[serde(default)]
pub vector_data: HashMap<VectorNameBuf, VectorDataConfig>,
#[serde(default)]
#[serde(skip_serializing_if = "HashMap::is_empty")]
pub sparse_vector_data: HashMap<VectorNameBuf, SparseVectorDataConfig>,
pub payload_storage_type: PayloadStorageType,
}
impl SegmentConfig {
pub fn quantization_config(&self, vector_name: &VectorName) -> Option<&QuantizationConfig> {
self.vector_data
.get(vector_name)
.and_then(|v| v.quantization_config.as_ref())
}
pub fn is_any_vector_indexed(&self) -> bool {
self.vector_data
.values()
.any(|config| config.index.is_indexed())
|| self
.sparse_vector_data
.values()
.any(|config| config.is_indexed())
}
pub fn is_any_on_disk(&self) -> bool {
self.vector_data
.values()
.any(|config| config.storage_type.is_on_disk())
|| self
.sparse_vector_data
.values()
.any(|config| config.index.index_type.is_on_disk())
}
pub fn is_appendable(&self) -> bool {
self.vector_data
.values()
.map(|vector_config| vector_config.is_appendable())
.chain(
self.sparse_vector_data
.values()
.map(|sparse_vector_config| {
sparse_vector_config.index.index_type.is_appendable()
}),
)
.all(|v| v)
}
pub fn check_compatible(&self, other: &Self) -> Result<(), String> {
let Self {
vector_data: _,
sparse_vector_data: _,
payload_storage_type: _,
} = self;
check_vectors_map_compatible(
&self.vector_data,
&other.vector_data,
VectorDataConfig::check_compatible,
)?;
check_vectors_map_compatible(
&self.sparse_vector_data,
&other.sparse_vector_data,
SparseVectorDataConfig::check_compatible,
)?;
Ok(())
}
}
fn check_vectors_map_compatible<C, F>(
this: &HashMap<String, C>,
other: &HashMap<String, C>,
check: F,
) -> Result<(), String>
where
F: Fn(&C, &C) -> Result<(), String>,
{
if this.len() != other.len() {
let expected_keys: Vec<String> = this.keys().map(|k| format!("{k:?}")).collect();
let actual_keys: Vec<String> = other.keys().map(|k| format!("{k:?}")).collect();
return Err(format!(
"Incompatible configs: expected vector storages with keys {expected_keys:?}, but got {actual_keys:?}"
));
}
for (vector_name, config) in this {
let Some(other_config) = other.get(vector_name) else {
return Err(format!(
"Incompatible configs: expected vector storage with key {vector_name:?} not found in other config"
));
};
check(config, other_config)
.map_err(|err| format!("Incompatible config for vector {vector_name:?}: {err}"))?;
}
Ok(())
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Eq, PartialEq, Copy, Clone)]
pub enum VectorStorageType {
Memory,
Mmap,
ChunkedMmap,
InRamChunkedMmap,
InRamMmap,
Empty,
}
#[cfg(any(test, feature = "testing"))]
#[allow(clippy::derivable_impls)]
impl Default for VectorStorageType {
fn default() -> Self {
VectorStorageType::InRamChunkedMmap
}
}
#[derive(
Default, Debug, Deserialize, Serialize, JsonSchema, Eq, PartialEq, Copy, Clone, Hash,
)]
#[serde(rename_all = "snake_case")]
pub enum VectorStorageDatatype {
#[default]
Float32,
Float16,
Uint8,
Turbo4,
}
#[derive(
Debug, Default, Deserialize, Serialize, JsonSchema, Eq, PartialEq, Copy, Clone, Hash,
)]
#[serde(rename_all = "snake_case")]
pub struct MultiVectorConfig {
pub comparator: MultiVectorComparator,
}
impl MultiVectorConfig {
fn check_compatible(&self, other: &Self) -> Result<(), String> {
let Self { comparator } = self;
if *comparator != other.comparator {
return Err(format!(
"Incompatible configs: expected multi-vector comparator {comparator:?}, but got {other_comparator:?}",
other_comparator = other.comparator
));
}
Ok(())
}
}
#[derive(
Debug, Default, Deserialize, Serialize, JsonSchema, Eq, PartialEq, Copy, Clone, Hash,
)]
#[serde(rename_all = "snake_case")]
pub enum MultiVectorComparator {
#[default]
MaxSim,
}
impl VectorStorageType {
pub fn from_on_disk(on_disk: bool) -> Self {
if on_disk {
Self::ChunkedMmap
} else {
Self::InRamChunkedMmap
}
}
pub fn is_on_disk(&self) -> bool {
match self {
Self::Memory | Self::InRamChunkedMmap | Self::InRamMmap => false,
Self::Mmap | Self::ChunkedMmap => true,
Self::Empty => true,
}
}
pub fn is_empty(&self) -> bool {
matches!(self, Self::Empty)
}
}
#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, JsonSchema, )]
#[serde(rename_all = "snake_case")]
pub struct VectorDataConfig {
pub size: usize,
pub distance: Distance,
pub storage_type: VectorStorageType,
pub index: Indexes,
pub quantization_config: Option<QuantizationConfig>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub multivector_config: Option<MultiVectorConfig>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub datatype: Option<VectorStorageDatatype>,
}
impl VectorDataConfig {
pub fn is_appendable(&self) -> bool {
let is_index_appendable = match self.index {
Indexes::Plain {} => true,
Indexes::Hnsw(_) => false,
};
let is_storage_appendable = match self.storage_type {
VectorStorageType::Memory => true,
VectorStorageType::Mmap => false,
VectorStorageType::ChunkedMmap => true,
VectorStorageType::InRamChunkedMmap => true,
VectorStorageType::InRamMmap => false,
VectorStorageType::Empty => false,
};
is_index_appendable && is_storage_appendable
}
pub fn check_compatible(&self, other: &Self) -> Result<(), String> {
let Self {
size,
distance,
storage_type: _,
index: _,
quantization_config: _,
multivector_config,
datatype,
} = self;
if *size != other.size {
return Err(format!(
"Incompatible configs: expected vector size {size}, but got {other_size}",
other_size = other.size
));
}
if *distance != other.distance {
return Err(format!(
"Incompatible configs: expected distance {distance:?}, but got {other_distance:?}",
other_distance = other.distance
));
}
let left_datatype = datatype.unwrap_or(VectorStorageDatatype::Float32);
let right_datatype = other.datatype.unwrap_or(VectorStorageDatatype::Float32);
if left_datatype != right_datatype {
return Err(format!(
"Incompatible configs: expected vector storage datatype {left_datatype:?}, but got {right_datatype:?}",
));
}
match (multivector_config, &other.multivector_config) {
(None, None) => {}
(Some(this), Some(other)) => {
MultiVectorConfig::check_compatible(this, other)?;
}
_ => {
return Err(format!(
"Incompatible configs: expected multivector config {this_multivector_config:?}, but got {other_multivector_config:?}",
this_multivector_config = multivector_config,
other_multivector_config = other.multivector_config
));
}
}
Ok(())
}
}
#[derive(
Copy, Clone, Debug, Default, Eq, PartialEq, Deserialize, Serialize, JsonSchema,
)]
#[serde(rename_all = "snake_case")]
pub enum SparseVectorStorageType {
#[default]
Mmap,
Empty,
}
impl SparseVectorStorageType {
pub fn is_on_disk(&self) -> bool {
match self {
Self::Mmap | Self::Empty => true,
}
}
}
#[derive(
Copy, Clone, Debug, PartialEq, Deserialize, Serialize, JsonSchema, Validate,
)]
#[serde(rename_all = "snake_case")]
pub struct SparseVectorDataConfig {
pub index: SparseIndexConfig,
#[serde(default = "default_sparse_vector_storage_type_when_not_in_config")]
pub storage_type: SparseVectorStorageType,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub modifier: Option<Modifier>,
}
fn default_sparse_vector_storage_type_when_not_in_config() -> SparseVectorStorageType {
SparseVectorStorageType::default()
}
impl SparseVectorDataConfig {
pub fn is_indexed(&self) -> bool {
true
}
pub fn check_compatible(&self, other: &Self) -> Result<(), String> {
let Self {
index: _,
storage_type: _,
modifier,
} = self;
if modifier != &other.modifier {
return Err(format!(
"Incompatible configs: expected sparse vector modifier {modifier:?}, but got {other_modifier:?}",
other_modifier = other.modifier
));
}
Ok(())
}
}
pub const DEFAULT_FULL_SCAN_THRESHOLD: usize = 10_000;
pub const DEFAULT_SPARSE_FULL_SCAN_THRESHOLD: usize = 5_000;
#[derive(Debug, Deserialize, Serialize, Clone)]
#[serde(rename_all = "snake_case")]
pub struct SegmentState {
#[serde(default)]
pub initial_version: Option<SeqNumberType>,
pub version: Option<SeqNumberType>,
pub config: SegmentConfig,
}
pub type RawGeoPoint = (f64, f64);
#[derive(
Debug,
Deserialize,
Serialize,
JsonSchema,
Clone,
Copy,
PartialEq,
Eq,
Hash,
Default,
PartialOrd,
Ord,
Pod,
Zeroable,
)]
#[serde(try_from = "GeoPointShadow")]
#[repr(C)]
pub struct GeoPoint {
pub lon: OrderedFloat<f64>,
pub lat: OrderedFloat<f64>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
pub struct GeoLineString {
pub points: Vec<GeoPoint>,
}
#[derive(Deserialize)]
struct GeoPointShadow {
pub lon: f64,
pub lat: f64,
}
#[derive(Debug)]
pub struct GeoPointValidationError {
pub lon: f64,
pub lat: f64,
}
impl std::fmt::Display for GeoPointValidationError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
formatter,
"Wrong format of GeoPoint payload: expected `lat` = {} within [-90;90] and `lon` = {} within [-180;180]",
self.lat, self.lon,
)
}
}
impl GeoPoint {
pub fn validate(lon: f64, lat: f64) -> Result<(), GeoPointValidationError> {
let max_lon = 180f64;
let min_lon = -180f64;
let max_lat = 90f64;
let min_lat = -90f64;
if !(min_lon..=max_lon).contains(&lon) || !(min_lat..=max_lat).contains(&lat) {
return Err(GeoPointValidationError { lon, lat });
}
Ok(())
}
pub fn new(lon: f64, lat: f64) -> Result<Self, GeoPointValidationError> {
Self::validate(lon, lat)?;
Ok(Self::new_unchecked(lon, lat))
}
pub const fn new_unchecked(lon: f64, lat: f64) -> Self {
GeoPoint {
lon: OrderedFloat(lon),
lat: OrderedFloat(lat),
}
}
}
impl TryFrom<GeoPointShadow> for GeoPoint {
type Error = GeoPointValidationError;
fn try_from(value: GeoPointShadow) -> Result<Self, Self::Error> {
let GeoPointShadow { lon, lat } = value;
GeoPoint::validate(lon, lat)?;
Ok(Self::new_unchecked(lon, lat))
}
}
impl From<GeoPoint> for geo::Point {
fn from(
GeoPoint {
lon: OrderedFloat(lon),
lat: OrderedFloat(lat),
}: GeoPoint,
) -> Self {
Self::new(lon, lat)
}
}
impl From<RawGeoPoint> for GeoPoint {
fn from((lon, lat): RawGeoPoint) -> Self {
GeoPoint::new(lon, lat).expect("invalid GeoPoint coordinates")
}
}
impl From<GeoPoint> for RawGeoPoint {
fn from(geo_point: GeoPoint) -> Self {
(geo_point.lon.0, geo_point.lat.0)
}
}
pub trait PayloadContainer {
fn get_value(&self, path: &JsonPath) -> MultiValue<&Value>;
fn get_value_cloned(&self, path: &JsonPath) -> MultiValue<Value> {
self.get_value(path).into_iter().cloned().collect()
}
}
macro_rules! payload_json {
($($tt:tt)*) => {
match ::serde_json::json!( { $($tt)* } ) {
::serde_json::Value::Object(map) => $crate::segment::types::Payload(map),
_ => unreachable!(),
}
};
}
#[allow(clippy::unnecessary_wraps)] fn payload_example() -> Option<Payload> {
Some(payload_json! {
"city": "London",
"color": "green",
})
}
#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize, JsonSchema, Hash)]
#[schemars(example = "payload_example")]
pub struct Payload(pub Map<String, Value>);
impl Payload {
pub fn merge(&mut self, value: &Payload) {
utils::merge_map(&mut self.0, &value.0)
}
pub fn merge_by_key(&mut self, value: &Payload, key: &JsonPath) {
JsonPath::value_set(Some(key), &mut self.0, &value.0);
}
pub fn remove(&mut self, path: &JsonPath) -> Vec<Value> {
path.value_remove(&mut self.0).to_vec()
}
pub fn len(&self) -> usize {
self.0.len()
}
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
pub fn contains_key(&self, key: &str) -> bool {
self.0.contains_key(key)
}
pub fn keys(&self) -> impl Iterator<Item = &String> {
self.0.keys()
}
}
impl PayloadContainer for Map<String, Value> {
fn get_value(&self, path: &JsonPath) -> MultiValue<&Value> {
path.value_get(self)
}
}
impl PayloadContainer for Payload {
fn get_value(&self, path: &JsonPath) -> MultiValue<&Value> {
path.value_get(&self.0)
}
}
impl PayloadContainer for OwnedPayloadRef<'_> {
fn get_value(&self, path: &JsonPath) -> MultiValue<&Value> {
path.value_get(self.as_ref())
}
}
impl Default for Payload {
fn default() -> Self {
Payload(Map::new())
}
}
impl IntoIterator for Payload {
type Item = (String, Value);
type IntoIter = serde_json::map::IntoIter;
fn into_iter(self) -> serde_json::map::IntoIter {
self.0.into_iter()
}
}
impl From<Map<String, Value>> for Payload {
fn from(value: serde_json::Map<String, Value>) -> Self {
Payload(value)
}
}
#[derive(Clone, Debug)]
pub enum OwnedPayloadRef<'a> {
Ref(&'a Map<String, Value>),
Owned(Rc<Map<String, Value>>),
}
impl Deref for OwnedPayloadRef<'_> {
type Target = Map<String, Value>;
fn deref(&self) -> &Self::Target {
match self {
OwnedPayloadRef::Ref(reference) => reference,
OwnedPayloadRef::Owned(owned) => owned.deref(),
}
}
}
impl AsRef<Map<String, Value>> for OwnedPayloadRef<'_> {
fn as_ref(&self) -> &Map<String, Value> {
match self {
OwnedPayloadRef::Ref(reference) => reference,
OwnedPayloadRef::Owned(owned) => owned.deref(),
}
}
}
impl From<Payload> for OwnedPayloadRef<'_> {
fn from(payload: Payload) -> Self {
OwnedPayloadRef::Owned(Rc::new(payload.0))
}
}
impl From<Map<String, Value>> for OwnedPayloadRef<'_> {
fn from(payload: Map<String, Value>) -> Self {
OwnedPayloadRef::Owned(Rc::new(payload))
}
}
impl<'a> From<&'a Payload> for OwnedPayloadRef<'a> {
fn from(payload: &'a Payload) -> Self {
OwnedPayloadRef::Ref(&payload.0)
}
}
impl<'a> From<&'a Map<String, Value>> for OwnedPayloadRef<'a> {
fn from(payload: &'a Map<String, Value>) -> Self {
OwnedPayloadRef::Ref(payload)
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq, Clone)]
#[serde(untagged, rename_all = "snake_case")]
pub enum PayloadVariant<T> {
List(Vec<T>),
Value(T),
}
#[derive(
Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Hash, Eq, EnumIter,
)]
#[serde(rename_all = "snake_case")]
pub enum PayloadSchemaType {
Keyword,
Integer,
Float,
Geo,
Text,
Bool,
Datetime,
Uuid,
}
impl PayloadSchemaType {
pub fn name(&self) -> &'static str {
serde_variant::to_variant_name(&self).unwrap_or("unknown")
}
pub fn expand(&self) -> PayloadSchemaParams {
match self {
Self::Keyword => PayloadSchemaParams::Keyword(KeywordIndexParams::default()),
Self::Integer => PayloadSchemaParams::Integer(IntegerIndexParams::default()),
Self::Float => PayloadSchemaParams::Float(FloatIndexParams::default()),
Self::Geo => PayloadSchemaParams::Geo(GeoIndexParams::default()),
Self::Text => PayloadSchemaParams::Text(TextIndexParams::default()),
Self::Bool => PayloadSchemaParams::Bool(BoolIndexParams::default()),
Self::Datetime => PayloadSchemaParams::Datetime(DatetimeIndexParams::default()),
Self::Uuid => PayloadSchemaParams::Uuid(UuidIndexParams::default()),
}
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Hash, Eq)]
#[serde(untagged, rename_all = "snake_case")]
pub enum PayloadSchemaParams {
Keyword(KeywordIndexParams),
Integer(IntegerIndexParams),
Float(FloatIndexParams),
Geo(GeoIndexParams),
Text(TextIndexParams),
Bool(BoolIndexParams),
Datetime(DatetimeIndexParams),
Uuid(UuidIndexParams),
}
impl PayloadSchemaParams {
pub fn name(&self) -> &'static str {
self.kind().name()
}
pub fn kind(&self) -> PayloadSchemaType {
match self {
PayloadSchemaParams::Keyword(_) => PayloadSchemaType::Keyword,
PayloadSchemaParams::Integer(_) => PayloadSchemaType::Integer,
PayloadSchemaParams::Float(_) => PayloadSchemaType::Float,
PayloadSchemaParams::Geo(_) => PayloadSchemaType::Geo,
PayloadSchemaParams::Text(_) => PayloadSchemaType::Text,
PayloadSchemaParams::Bool(_) => PayloadSchemaType::Bool,
PayloadSchemaParams::Datetime(_) => PayloadSchemaType::Datetime,
PayloadSchemaParams::Uuid(_) => PayloadSchemaType::Uuid,
}
}
pub fn tenant_optimization(&self) -> bool {
match self {
PayloadSchemaParams::Keyword(keyword) => keyword.is_tenant.unwrap_or_default(),
PayloadSchemaParams::Integer(integer) => integer.is_principal.unwrap_or_default(),
PayloadSchemaParams::Float(float) => float.is_principal.unwrap_or_default(),
PayloadSchemaParams::Datetime(datetime) => datetime.is_principal.unwrap_or_default(),
PayloadSchemaParams::Uuid(uuid) => uuid.is_tenant.unwrap_or_default(),
PayloadSchemaParams::Geo(_)
| PayloadSchemaParams::Text(_)
| PayloadSchemaParams::Bool(_) => false,
}
}
pub fn is_on_disk(&self) -> bool {
match self {
PayloadSchemaParams::Keyword(i) => i.on_disk.unwrap_or_default(),
PayloadSchemaParams::Integer(i) => i.on_disk.unwrap_or_default(),
PayloadSchemaParams::Float(i) => i.on_disk.unwrap_or_default(),
PayloadSchemaParams::Datetime(i) => i.on_disk.unwrap_or_default(),
PayloadSchemaParams::Uuid(i) => i.on_disk.unwrap_or_default(),
PayloadSchemaParams::Text(i) => i.on_disk.unwrap_or_default(),
PayloadSchemaParams::Geo(i) => i.on_disk.unwrap_or_default(),
PayloadSchemaParams::Bool(i) => i.on_disk.unwrap_or_default(),
}
}
pub fn enable_hnsw(&self) -> bool {
match self {
PayloadSchemaParams::Keyword(params) => params.enable_hnsw.unwrap_or(true),
PayloadSchemaParams::Integer(params) => params.enable_hnsw.unwrap_or(true),
PayloadSchemaParams::Float(params) => params.enable_hnsw.unwrap_or(true),
PayloadSchemaParams::Datetime(params) => params.enable_hnsw.unwrap_or(true),
PayloadSchemaParams::Uuid(params) => params.enable_hnsw.unwrap_or(true),
PayloadSchemaParams::Text(params) => params.enable_hnsw.unwrap_or(true),
PayloadSchemaParams::Geo(params) => params.enable_hnsw.unwrap_or(true),
PayloadSchemaParams::Bool(params) => params.enable_hnsw.unwrap_or(true),
}
}
}
impl Validate for PayloadSchemaParams {
fn validate(&self) -> Result<(), ValidationErrors> {
match self {
PayloadSchemaParams::Keyword(_) => Ok(()),
PayloadSchemaParams::Integer(integer_index_params) => integer_index_params.validate(),
PayloadSchemaParams::Float(_) => Ok(()),
PayloadSchemaParams::Geo(_) => Ok(()),
PayloadSchemaParams::Text(_) => Ok(()),
PayloadSchemaParams::Bool(_) => Ok(()),
PayloadSchemaParams::Datetime(_) => Ok(()),
PayloadSchemaParams::Uuid(_) => Ok(()),
}
}
}
#[derive(Clone, Debug, Eq, Deserialize, Serialize, JsonSchema)]
#[serde(untagged, rename_all = "snake_case")]
pub enum PayloadFieldSchema {
FieldType(PayloadSchemaType),
FieldParams(PayloadSchemaParams),
}
impl PartialEq for PayloadFieldSchema {
fn eq(&self, other: &Self) -> bool {
match (self, other) {
(Self::FieldType(this), Self::FieldType(other)) => this == other,
(Self::FieldParams(this), Self::FieldParams(other)) => this == other,
(Self::FieldType(this), Self::FieldParams(other)) => &this.expand() == other,
(Self::FieldParams(this), Self::FieldType(other)) => this == &other.expand(),
}
}
}
impl hash::Hash for PayloadFieldSchema {
fn hash<H: hash::Hasher>(&self, state: &mut H) {
match self {
PayloadFieldSchema::FieldType(default) => default.expand().hash(state),
PayloadFieldSchema::FieldParams(params) => params.hash(state),
}
}
}
impl Validate for PayloadFieldSchema {
fn validate(&self) -> Result<(), ValidationErrors> {
match self {
PayloadFieldSchema::FieldType(_) => Ok(()), PayloadFieldSchema::FieldParams(payload_schema_params) => {
payload_schema_params.validate()
}
}
}
}
impl Display for PayloadFieldSchema {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
match self {
PayloadFieldSchema::FieldType(t) => write!(f, "{}", t.name()),
PayloadFieldSchema::FieldParams(params) => match params {
PayloadSchemaParams::Keyword(_)
| PayloadSchemaParams::Float(_)
| PayloadSchemaParams::Geo(_)
| PayloadSchemaParams::Bool(_)
| PayloadSchemaParams::Datetime(_)
| PayloadSchemaParams::Uuid(_) => write!(f, "{}", params.name()),
PayloadSchemaParams::Integer(integer_params) => {
let range = integer_params.range.unwrap_or(true);
let lookup = integer_params.lookup.unwrap_or(true);
if range && lookup {
write!(f, "integer")
} else {
write!(f, "integer (with range: {range}, lookup: {lookup})")
}
}
PayloadSchemaParams::Text(text_params) => {
if text_params.phrase_matching.unwrap_or_default() {
write!(f, "text (with phrase_matching: true)")
} else {
write!(f, "text")
}
}
},
}
}
}
impl PayloadFieldSchema {
pub fn expand(&self) -> Cow<'_, PayloadSchemaParams> {
match self {
PayloadFieldSchema::FieldType(t) => Cow::Owned(t.expand()),
PayloadFieldSchema::FieldParams(p) => Cow::Borrowed(p),
}
}
pub fn name(&self) -> &'static str {
match self {
PayloadFieldSchema::FieldType(field_type) => field_type.name(),
PayloadFieldSchema::FieldParams(field_params) => field_params.name(),
}
}
pub fn is_tenant(&self) -> bool {
match self {
PayloadFieldSchema::FieldType(_) => false,
PayloadFieldSchema::FieldParams(params) => params.tenant_optimization(),
}
}
pub fn is_on_disk(&self) -> bool {
match self {
PayloadFieldSchema::FieldType(_) => false,
PayloadFieldSchema::FieldParams(params) => params.is_on_disk(),
}
}
pub fn kind(&self) -> PayloadSchemaType {
match self {
PayloadFieldSchema::FieldType(t) => *t,
PayloadFieldSchema::FieldParams(p) => p.kind(),
}
}
pub fn supports_match(&self) -> bool {
match self {
PayloadFieldSchema::FieldType(payload_schema_type) => match payload_schema_type {
PayloadSchemaType::Keyword => true,
PayloadSchemaType::Integer => true,
PayloadSchemaType::Uuid => true,
PayloadSchemaType::Bool => true,
PayloadSchemaType::Float => false,
PayloadSchemaType::Geo => false,
PayloadSchemaType::Text => false,
PayloadSchemaType::Datetime => false,
},
PayloadFieldSchema::FieldParams(payload_schema_params) => match payload_schema_params {
PayloadSchemaParams::Keyword(_) => true,
PayloadSchemaParams::Integer(integer_index_params) => {
integer_index_params.lookup == Some(true)
}
PayloadSchemaParams::Uuid(_) => true,
PayloadSchemaParams::Bool(_) => true,
PayloadSchemaParams::Float(_) => false,
PayloadSchemaParams::Geo(_) => false,
PayloadSchemaParams::Text(_) => false,
PayloadSchemaParams::Datetime(_) => false,
},
}
}
pub fn enable_hnsw(&self) -> bool {
match self {
PayloadFieldSchema::FieldType(_) => true,
PayloadFieldSchema::FieldParams(p) => p.enable_hnsw(),
}
}
}
impl From<PayloadSchemaType> for PayloadFieldSchema {
fn from(payload_schema_type: PayloadSchemaType) -> Self {
PayloadFieldSchema::FieldType(payload_schema_type)
}
}
impl TryFrom<PayloadIndexInfo> for PayloadFieldSchema {
type Error = String;
fn try_from(index_info: PayloadIndexInfo) -> Result<Self, Self::Error> {
let PayloadIndexInfo {
data_type,
params,
points: _,
} = index_info;
match params {
None => Ok(PayloadFieldSchema::FieldType(data_type)),
Some(params) if data_type == params.kind() => {
Ok(PayloadFieldSchema::FieldParams(params))
}
Some(params) => Err(format!(
"payload field with type {data_type:?} has parameters of type {:?}",
params.kind(),
)),
}
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
#[serde(untagged)]
pub enum ValueVariants {
String(String),
Integer(IntPayloadType),
Bool(bool),
}
impl ValueVariants {
pub fn to_value(&self) -> Value {
match self {
ValueVariants::String(keyword) => Value::String(keyword.clone()),
&ValueVariants::Integer(integer) => Value::Number(integer.into()),
&ValueVariants::Bool(flag) => Value::Bool(flag),
}
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq)]
#[serde(untagged)]
pub enum AnyVariants {
Strings(IndexSet<String, FnvBuildHasher>),
Integers(IndexSet<IntPayloadType, FnvBuildHasher>),
}
impl Hash for AnyVariants {
fn hash<H: Hasher>(&self, state: &mut H) {
mem::discriminant(self).hash(state);
match self {
AnyVariants::Strings(index_set) => {
for item in index_set {
item.hash(state);
}
}
AnyVariants::Integers(index_set) => {
for item in index_set {
item.hash(state);
}
}
}
}
}
impl AnyVariants {
pub fn len(&self) -> usize {
match self {
AnyVariants::Strings(index_set) => index_set.len(),
AnyVariants::Integers(index_set) => index_set.len(),
}
}
pub fn is_empty(&self) -> bool {
match self {
AnyVariants::Strings(index_set) => index_set.is_empty(),
AnyVariants::Integers(index_set) => index_set.is_empty(),
}
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
#[serde(rename_all = "snake_case")]
pub struct MatchValue {
pub value: ValueVariants,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
#[serde(rename_all = "snake_case")]
pub struct MatchText {
pub text: String,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
#[serde(rename_all = "snake_case")]
pub struct MatchTextAny {
pub text_any: String,
}
impl<S: Into<String>> From<S> for MatchText {
fn from(text: S) -> Self {
MatchText { text: text.into() }
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
#[serde(rename_all = "snake_case")]
pub struct MatchPhrase {
pub phrase: String,
}
impl<S: Into<String>> From<S> for MatchPhrase {
fn from(text: S) -> Self {
MatchPhrase {
phrase: text.into(),
}
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
#[serde(rename_all = "snake_case")]
pub struct MatchAny {
pub any: AnyVariants,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
#[serde(rename_all = "snake_case")]
pub struct MatchExcept {
pub except: AnyVariants,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq)]
#[serde(untagged, rename_all = "snake_case")]
pub enum MatchInterface {
Value(MatchValue),
Text(MatchText),
TextAny(MatchTextAny),
Phrase(MatchPhrase),
Any(MatchAny),
Except(MatchExcept),
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
#[serde(untagged, from = "MatchInterface")]
pub enum Match {
Value(MatchValue),
Text(MatchText),
TextAny(MatchTextAny),
Phrase(MatchPhrase),
Any(MatchAny),
Except(MatchExcept),
}
impl Match {
pub fn new_value(value: ValueVariants) -> Self {
Self::Value(MatchValue { value })
}
pub fn new_text(text: &str) -> Self {
Self::Text(MatchText { text: text.into() })
}
pub fn new_any(any: AnyVariants) -> Self {
Self::Any(MatchAny { any })
}
pub fn new_except(except: AnyVariants) -> Self {
Self::Except(MatchExcept { except })
}
}
impl From<AnyVariants> for Match {
fn from(any: AnyVariants) -> Self {
Self::Any(MatchAny { any })
}
}
impl From<MatchInterface> for Match {
fn from(value: MatchInterface) -> Self {
match value {
MatchInterface::Value(value) => Self::Value(MatchValue { value: value.value }),
MatchInterface::Text(text) => Self::Text(MatchText { text: text.text }),
MatchInterface::TextAny(text_any) => Self::TextAny(MatchTextAny {
text_any: text_any.text_any,
}),
MatchInterface::Any(any) => Self::Any(MatchAny { any: any.any }),
MatchInterface::Except(except) => Self::Except(MatchExcept {
except: except.except,
}),
MatchInterface::Phrase(MatchPhrase { phrase }) => Self::Phrase(MatchPhrase { phrase }),
}
}
}
impl From<bool> for Match {
fn from(flag: bool) -> Self {
Self::Value(MatchValue {
value: ValueVariants::Bool(flag),
})
}
}
impl From<String> for Match {
fn from(keyword: String) -> Self {
Self::Value(MatchValue {
value: ValueVariants::String(keyword),
})
}
}
impl From<EcoString> for Match {
fn from(keyword: EcoString) -> Self {
Self::Value(MatchValue {
value: ValueVariants::String(keyword.into()),
})
}
}
impl From<IntPayloadType> for Match {
fn from(integer: IntPayloadType) -> Self {
Self::Value(MatchValue {
value: ValueVariants::Integer(integer),
})
}
}
impl From<Vec<String>> for Match {
fn from(keywords: Vec<String>) -> Self {
let keywords: IndexSet<String, FnvBuildHasher> = keywords.into_iter().collect();
Self::Any(MatchAny {
any: AnyVariants::Strings(keywords),
})
}
}
impl From<ValueVariants> for Match {
fn from(value: ValueVariants) -> Self {
Self::Value(MatchValue { value })
}
}
impl From<Vec<String>> for MatchExcept {
fn from(keywords: Vec<String>) -> Self {
let keywords: IndexSet<String, FnvBuildHasher> = keywords.into_iter().collect();
MatchExcept {
except: AnyVariants::Strings(keywords),
}
}
}
impl From<Vec<IntPayloadType>> for Match {
fn from(integers: Vec<IntPayloadType>) -> Self {
let integers: IndexSet<_, FnvBuildHasher> = integers.into_iter().collect();
Self::Any(MatchAny {
any: AnyVariants::Integers(integers),
})
}
}
impl From<Vec<IntPayloadType>> for MatchExcept {
fn from(integers: Vec<IntPayloadType>) -> Self {
let integers: IndexSet<_, FnvBuildHasher> = integers.into_iter().collect();
MatchExcept {
except: AnyVariants::Integers(integers),
}
}
}
#[derive(Copy, Clone, Debug, Eq, PartialEq, Serialize, JsonSchema)]
#[serde(untagged)]
pub enum RangeInterface {
Float(Range<OrderedFloat<FloatPayloadType>>),
DateTime(Range<DateTimePayloadType>),
}
impl Hash for RangeInterface {
fn hash<H: hash::Hasher>(&self, state: &mut H) {
match self {
RangeInterface::Float(range) => {
let Range { lt, gt, gte, lte } = range;
lt.hash(state);
gt.hash(state);
gte.hash(state);
lte.hash(state);
}
RangeInterface::DateTime(range) => {
let Range { lt, gt, gte, lte } = range;
lt.hash(state);
gt.hash(state);
gte.hash(state);
lte.hash(state);
}
}
}
}
#[derive(serde::Deserialize)]
#[serde(untagged)]
enum RangeInterfaceUntagged {
Float(Range<OrderedFloatPayloadType>),
DateTime(Range<DateTimePayloadType>),
}
impl<'de> serde::Deserialize<'de> for RangeInterface {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
if !deserializer.is_human_readable() {
return RangeInterfaceUntagged::deserialize(deserializer).map(|parsed| match parsed {
RangeInterfaceUntagged::Float(r) => RangeInterface::Float(r),
RangeInterfaceUntagged::DateTime(r) => RangeInterface::DateTime(r),
});
}
let value = serde_json::Value::deserialize(deserializer)?;
if let Some(obj) = value.as_object() {
let keys = ["lt", "gt", "lte", "gte"];
let has_string_bound = keys
.iter()
.any(|k| obj.get(*k).is_some_and(|v| v.is_string()));
if has_string_bound {
return serde_json::from_value::<Range<DateTimePayloadType>>(value)
.map(RangeInterface::DateTime)
.map_err(serde::de::Error::custom);
}
}
let parsed = serde_json::from_value::<RangeInterfaceUntagged>(value)
.map_err(serde::de::Error::custom)?;
Ok(match parsed {
RangeInterfaceUntagged::Float(r) => RangeInterface::Float(r),
RangeInterfaceUntagged::DateTime(r) => RangeInterface::DateTime(r),
})
}
}
type OrderedFloatPayloadType = OrderedFloat<FloatPayloadType>;
#[macro_rules_attribute::macro_rules_derive(crate::segment::common::macros::schemars_rename_generics)]
#[derive_args(< OrderedFloatPayloadType > => "Range", < DateTimePayloadType > => "DatetimeRange")]
#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Deserialize, Serialize)]
#[serde(rename_all = "snake_case")]
pub struct Range<T> {
pub lt: Option<T>,
pub gt: Option<T>,
pub gte: Option<T>,
pub lte: Option<T>,
}
impl<T: Copy> Range<T> {
pub fn map<U, F: Fn(T) -> U>(&self, f: F) -> Range<U> {
let Self { lt, gt, gte, lte } = self;
Range {
lt: lt.map(&f),
gt: gt.map(&f),
gte: gte.map(&f),
lte: lte.map(&f),
}
}
}
impl<T: Copy + PartialOrd> Range<T> {
pub fn check_range(&self, number: T) -> bool {
let Self { lt, gt, gte, lte } = self;
lt.is_none_or(|x| number < x)
&& gt.is_none_or(|x| number > x)
&& lte.is_none_or(|x| number <= x)
&& gte.is_none_or(|x| number >= x)
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Copy, Clone, PartialEq, Eq, Hash)]
#[serde(rename_all = "snake_case")]
pub struct ValuesCount {
pub lt: Option<usize>,
pub gt: Option<usize>,
pub gte: Option<usize>,
pub lte: Option<usize>,
}
impl ValuesCount {
pub fn check_count(&self, count: usize) -> bool {
let Self { lt, gt, gte, lte } = self;
lt.is_none_or(|x| count < x)
&& gt.is_none_or(|x| count > x)
&& lte.is_none_or(|x| count <= x)
&& gte.is_none_or(|x| count >= x)
}
pub fn check_count_from(&self, value: &Value) -> bool {
let count = match value {
Value::Null => 0,
Value::Array(array) => array.len(),
Value::Bool(_) | Value::Number(_) | Value::String(_) | Value::Object(_) => 1,
};
self.check_count(count)
}
}
#[cfg(test)]
impl From<std::ops::Range<usize>> for ValuesCount {
fn from(range: std::ops::Range<usize>) -> Self {
Self {
gte: Some(range.start),
lt: Some(range.end),
gt: None,
lte: None,
}
}
}
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Deserialize, Serialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub struct GeoBoundingBox {
pub top_left: GeoPoint,
pub bottom_right: GeoPoint,
}
impl GeoBoundingBox {
pub fn check_point(&self, point: &GeoPoint) -> bool {
let longitude_check = if self.top_left.lon > self.bottom_right.lon {
point.lon > self.top_left.lon || point.lon < self.bottom_right.lon
} else {
self.top_left.lon < point.lon && point.lon < self.bottom_right.lon
};
let latitude_check = self.bottom_right.lat < point.lat && point.lat < self.top_left.lat;
longitude_check && latitude_check
}
}
#[derive(Copy, Clone, Debug, Eq, PartialEq, Deserialize, Serialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub struct GeoRadius {
pub center: GeoPoint,
pub radius: OrderedFloat<f64>,
}
impl Hash for GeoRadius {
fn hash<H: hash::Hasher>(&self, state: &mut H) {
let GeoRadius { center, radius } = self;
center.hash(state);
OrderedFloat(*radius).hash(state);
}
}
impl GeoRadius {
pub fn check_point(&self, point: &GeoPoint) -> bool {
let query_center = Point::from(self.center);
Haversine.distance(query_center, Point::from(*point)) < self.radius.0
}
}
#[derive(Deserialize)]
pub struct GeoPolygonShadow {
pub exterior: GeoLineString,
pub interiors: Option<Vec<GeoLineString>>,
}
pub struct PolygonWrapper {
pub polygon: Polygon,
}
impl PolygonWrapper {
pub fn check_point(&self, point: &GeoPoint) -> bool {
let point_new = Point::new(point.lon.0, point.lat.0);
self.polygon.contains(&point_new)
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
#[serde(try_from = "GeoPolygonShadow", rename_all = "snake_case")]
pub struct GeoPolygon {
pub exterior: GeoLineString,
pub interiors: Option<Vec<GeoLineString>>,
}
impl GeoPolygon {
pub fn validate_line_string(line: &GeoLineString) -> OperationResult<()> {
if line.points.len() <= 3 {
return Err(OperationError::validation_error(format!(
"polygon invalid, the size must be at least 4, got {}",
line.points.len()
)));
}
if let (Some(first), Some(last)) = (line.points.first(), line.points.last())
&& ((first.lat - last.lat).abs() > f64::EPSILON
|| (first.lon - last.lon).abs() > f64::EPSILON)
{
return Err(OperationError::validation_error(
"polygon invalid, the first and the last points should be the same to form a closed line",
));
}
Ok(())
}
pub fn convert(&self) -> PolygonWrapper {
let exterior_line: LineString = LineString(
self.exterior
.points
.iter()
.map(|p| Coord {
x: p.lon.0,
y: p.lat.0,
})
.collect(),
);
let interior_lines: Vec<LineString> = match &self.interiors {
None => vec![],
Some(interiors) => interiors
.iter()
.map(|interior_points| {
interior_points
.points
.iter()
.map(|p| Coord {
x: p.lon.0,
y: p.lat.0,
})
.collect()
})
.map(LineString)
.collect(),
};
PolygonWrapper {
polygon: Polygon::new(exterior_line, interior_lines),
}
}
}
impl TryFrom<GeoPolygonShadow> for GeoPolygon {
type Error = OperationError;
fn try_from(value: GeoPolygonShadow) -> OperationResult<Self> {
let GeoPolygonShadow {
exterior,
interiors,
} = value;
Self::validate_line_string(&exterior)?;
if let Some(interiors) = &interiors {
for interior in interiors {
Self::validate_line_string(interior)?;
}
}
Ok(GeoPolygon {
exterior,
interiors,
})
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Validate, Clone, PartialEq, Eq, Hash)]
#[validate(schema(function = "validate_field_condition"))]
#[serde(rename_all = "snake_case")]
pub struct FieldCondition {
pub key: PayloadKeyType,
#[serde(skip_serializing_if = "Option::is_none")]
pub r#match: Option<Match>,
#[serde(skip_serializing_if = "Option::is_none")]
pub range: Option<RangeInterface>,
#[serde(skip_serializing_if = "Option::is_none")]
pub geo_bounding_box: Option<GeoBoundingBox>,
#[serde(skip_serializing_if = "Option::is_none")]
pub geo_radius: Option<GeoRadius>,
#[serde(skip_serializing_if = "Option::is_none")]
pub geo_polygon: Option<GeoPolygon>,
#[serde(skip_serializing_if = "Option::is_none")]
pub values_count: Option<ValuesCount>,
#[serde(skip_serializing_if = "Option::is_none")]
pub is_empty: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub is_null: Option<bool>,
}
impl FieldCondition {
pub fn new_match(key: PayloadKeyType, r#match: Match) -> Self {
Self {
key,
r#match: Some(r#match),
range: None,
geo_bounding_box: None,
geo_radius: None,
geo_polygon: None,
values_count: None,
is_empty: None,
is_null: None,
}
}
pub fn new_range(key: PayloadKeyType, range: Range<OrderedFloat<FloatPayloadType>>) -> Self {
Self {
key,
r#match: None,
range: Some(RangeInterface::Float(range)),
geo_bounding_box: None,
geo_radius: None,
geo_polygon: None,
values_count: None,
is_empty: None,
is_null: None,
}
}
pub fn new_datetime_range(
key: PayloadKeyType,
datetime_range: Range<DateTimePayloadType>,
) -> Self {
Self {
key,
r#match: None,
range: Some(RangeInterface::DateTime(datetime_range)),
geo_bounding_box: None,
geo_radius: None,
geo_polygon: None,
values_count: None,
is_empty: None,
is_null: None,
}
}
pub fn new_geo_bounding_box(key: PayloadKeyType, geo_bounding_box: GeoBoundingBox) -> Self {
Self {
key,
r#match: None,
range: None,
geo_bounding_box: Some(geo_bounding_box),
geo_radius: None,
geo_polygon: None,
values_count: None,
is_empty: None,
is_null: None,
}
}
pub fn new_geo_radius(key: PayloadKeyType, geo_radius: GeoRadius) -> Self {
Self {
key,
r#match: None,
range: None,
geo_bounding_box: None,
geo_radius: Some(geo_radius),
geo_polygon: None,
values_count: None,
is_empty: None,
is_null: None,
}
}
pub fn new_geo_polygon(key: PayloadKeyType, geo_polygon: GeoPolygon) -> Self {
Self {
key,
r#match: None,
range: None,
geo_bounding_box: None,
geo_radius: None,
geo_polygon: Some(geo_polygon),
values_count: None,
is_empty: None,
is_null: None,
}
}
pub fn new_values_count(key: PayloadKeyType, values_count: ValuesCount) -> Self {
Self {
key,
r#match: None,
range: None,
geo_bounding_box: None,
geo_radius: None,
geo_polygon: None,
values_count: Some(values_count),
is_empty: None,
is_null: None,
}
}
pub fn new_is_empty(key: PayloadKeyType, is_empty: bool) -> Self {
Self {
key,
r#match: None,
range: None,
geo_bounding_box: None,
geo_radius: None,
geo_polygon: None,
values_count: None,
is_empty: Some(is_empty),
is_null: None,
}
}
pub fn new_is_null(key: PayloadKeyType, is_null: bool) -> Self {
Self {
key,
r#match: None,
range: None,
geo_bounding_box: None,
geo_radius: None,
geo_polygon: None,
values_count: None,
is_empty: None,
is_null: Some(is_null),
}
}
pub fn all_fields_none(&self) -> bool {
matches!(
self,
FieldCondition {
r#match: None,
range: None,
geo_bounding_box: None,
geo_radius: None,
geo_polygon: None,
values_count: None,
key: _,
is_empty: None,
is_null: None,
}
)
}
fn input_size(&self) -> usize {
if self.r#match.is_none() {
return 0;
}
match self.r#match.as_ref().unwrap() {
Match::Any(match_any) => match_any.any.len(),
Match::Except(match_except) => match_except.except.len(),
Match::Value(_) => 0,
Match::Text(_) => 0,
Match::Phrase(_) => 0,
Match::TextAny(_) => 0,
}
}
}
pub fn validate_field_condition(field_condition: &FieldCondition) -> Result<(), ValidationError> {
if field_condition.all_fields_none() {
Err(ValidationError::new(
"At least one field condition must be specified",
))
} else {
Ok(())
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
pub struct PayloadField {
pub key: PayloadKeyType,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
pub struct IsEmptyCondition {
pub is_empty: PayloadField,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
pub struct IsNullCondition {
pub is_null: PayloadField,
}
impl From<JsonPath> for IsNullCondition {
fn from(key: PayloadKeyType) -> Self {
IsNullCondition {
is_null: PayloadField { key },
}
}
}
impl From<JsonPath> for IsEmptyCondition {
fn from(key: PayloadKeyType) -> Self {
IsEmptyCondition {
is_empty: PayloadField { key },
}
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq)]
pub struct HasIdCondition {
#[schemars(schema_with = "HashSet::<PointIdType>::json_schema")]
pub has_id: MaybeArc<AHashSet<PointIdType>>,
}
impl Hash for HasIdCondition {
fn hash<H: hash::Hasher>(&self, state: &mut H) {
unordered_hash_unique(state, self.has_id.iter());
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
pub struct HasVectorCondition {
pub has_vector: VectorNameBuf,
}
impl From<VectorNameBuf> for HasVectorCondition {
fn from(vector: VectorNameBuf) -> Self {
HasVectorCondition { has_vector: vector }
}
}
const HAS_ID_CONDITION_ARC_THRESHOLD: usize = 1_000;
impl From<AHashSet<PointIdType>> for HasIdCondition {
fn from(has_id: AHashSet<PointIdType>) -> Self {
if has_id.len() > HAS_ID_CONDITION_ARC_THRESHOLD {
HasIdCondition {
has_id: MaybeArc::arc(has_id),
}
} else {
HasIdCondition {
has_id: MaybeArc::no_arc(has_id),
}
}
}
}
impl FromIterator<PointIdType> for HasIdCondition {
fn from_iter<T: IntoIterator<Item = PointIdType>>(iter: T) -> Self {
let items: AHashSet<_> = iter.into_iter().collect();
Self::from(items)
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Validate, Hash)]
pub struct Nested {
pub key: PayloadKeyType,
#[validate(nested)]
pub filter: Filter,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Validate, Hash)]
pub struct NestedCondition {
#[validate(nested)]
pub nested: Nested,
}
impl NestedCondition {
pub fn new(nested: Nested) -> Self {
Self { nested }
}
pub fn raw_key(&self) -> &PayloadKeyType {
&self.nested.key
}
pub fn array_key(&self) -> PayloadKeyType {
self.raw_key().array_key()
}
pub fn filter(&self) -> &Filter {
&self.nested.filter
}
}
#[derive(Clone, Debug, Serialize, JsonSchema, PartialEq, Eq, Hash)]
#[serde(untagged)]
#[serde(
expecting = "Expected some form of condition, which can be a field condition (like {\"key\": ..., \"match\": ... }), or some other mentioned in the documentation: https://qdrant.tech/documentation/concepts/filtering/#filtering-conditions"
)]
#[allow(clippy::large_enum_variant)]
pub enum Condition {
Field(FieldCondition),
IsEmpty(IsEmptyCondition),
IsNull(IsNullCondition),
HasId(HasIdCondition),
HasVector(HasVectorCondition),
Nested(NestedCondition),
Filter(Filter),
#[serde(skip)]
CustomIdChecker(CustomIdChecker),
}
#[derive(Deserialize)]
#[serde(untagged)]
#[serde(
expecting = "Expected some form of condition, which can be a field condition (like {\"key\": ..., \"match\": ... }), or some other mentioned in the documentation: https://qdrant.tech/documentation/concepts/filtering/#filtering-conditions"
)]
#[allow(clippy::large_enum_variant, dead_code)]
enum ConditionUntagged {
Field(FieldCondition),
IsEmpty(IsEmptyCondition),
IsNull(IsNullCondition),
HasId(HasIdCondition),
HasVector(HasVectorCondition),
Nested(NestedCondition),
Filter(Filter),
#[serde(skip)]
CustomIdChecker(CustomIdChecker),
}
impl From<ConditionUntagged> for Condition {
fn from(condition: ConditionUntagged) -> Self {
match condition {
ConditionUntagged::Field(condition) => Condition::Field(condition),
ConditionUntagged::IsEmpty(condition) => Condition::IsEmpty(condition),
ConditionUntagged::IsNull(condition) => Condition::IsNull(condition),
ConditionUntagged::HasId(condition) => Condition::HasId(condition),
ConditionUntagged::HasVector(condition) => Condition::HasVector(condition),
ConditionUntagged::Nested(condition) => Condition::Nested(condition),
ConditionUntagged::Filter(condition) => Condition::Filter(condition),
ConditionUntagged::CustomIdChecker(condition) => Condition::CustomIdChecker(condition),
}
}
}
impl<'de> serde::Deserialize<'de> for Condition {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let value = serde_value::Value::deserialize(deserializer)?;
if let serde_value::Value::Map(obj) = &value
&& obj.contains_key(&serde_value::Value::String("key".into()))
{
return value
.deserialize_into()
.map(Condition::Field)
.map_err(serde::de::Error::custom);
}
value
.deserialize_into::<ConditionUntagged>()
.map(Condition::from)
.map_err(serde::de::Error::custom)
}
}
impl Condition {
pub fn new_custom(checker: Arc<dyn CustomIdCheckerCondition + Send + Sync + 'static>) -> Self {
Condition::CustomIdChecker(CustomIdChecker(checker))
}
}
#[derive(Debug, Clone)]
pub struct CustomIdChecker(pub Arc<dyn CustomIdCheckerCondition + Send + Sync + 'static>);
impl Hash for CustomIdChecker {
fn hash<H: hash::Hasher>(&self, state: &mut H) {
std::ptr::hash(Arc::as_ptr(&self.0), state);
}
}
impl PartialEq for CustomIdChecker {
fn eq(&self, other: &Self) -> bool {
Arc::ptr_eq(&self.0, &other.0)
}
}
impl Eq for CustomIdChecker {}
impl Condition {
pub fn new_nested(key: JsonPath, filter: Filter) -> Self {
Self::Nested(NestedCondition {
nested: Nested { key, filter },
})
}
pub fn size_estimation(&self) -> usize {
match self {
Condition::Field(field_condition) => field_condition.input_size(),
Condition::HasId(has_id_condition) => has_id_condition.has_id.len(),
Condition::Filter(filter) => filter.max_condition_input_size(),
Condition::Nested(nested) => nested.filter().max_condition_input_size(),
Condition::IsEmpty(_)
| Condition::IsNull(_)
| Condition::HasVector(_)
| Condition::CustomIdChecker(_) => 0,
}
}
pub fn sub_conditions_count(&self) -> usize {
match self {
Condition::Nested(nested_condition) => {
nested_condition.filter().total_conditions_count()
}
Condition::Filter(filter) => filter.total_conditions_count(),
Condition::Field(_)
| Condition::IsEmpty(_)
| Condition::IsNull(_)
| Condition::CustomIdChecker(_)
| Condition::HasId(_)
| Condition::HasVector(_) => 1,
}
}
pub fn targeted_key(&self) -> Option<PayloadKeyType> {
match self {
Condition::Field(field_condition) => Some(field_condition.key.clone()),
Condition::IsEmpty(is_empty_condition) => Some(is_empty_condition.is_empty.key.clone()),
Condition::IsNull(is_null_condition) => Some(is_null_condition.is_null.key.clone()),
Condition::Nested(nested_condition) => Some(nested_condition.array_key()),
Condition::Filter(filter) => filter.iter_conditions().find_map(|c| c.targeted_key()),
Condition::HasId(_) | Condition::HasVector(_) | Condition::CustomIdChecker(_) => None,
}
}
}
impl Validate for Condition {
fn validate(&self) -> Result<(), ValidationErrors> {
match self {
Condition::HasId(_)
| Condition::IsEmpty(_)
| Condition::IsNull(_)
| Condition::HasVector(_) => Ok(()),
Condition::Field(field_condition) => field_condition.validate(),
Condition::Nested(nested_condition) => nested_condition.validate(),
Condition::Filter(filter) => filter.validate(),
Condition::CustomIdChecker(_) => Ok(()),
}
}
}
pub trait CustomIdCheckerCondition: fmt::Debug {
fn estimate_cardinality(&self, points: usize) -> CardinalityEstimation;
fn check(&self, point_id: ExtendedPointId) -> bool;
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Hash)]
#[serde(untagged, rename_all = "snake_case")]
#[serde(
expecting = "Expected a boolean, an array of strings, or an object with an include/exclude field"
)]
pub enum WithPayloadInterface {
Bool(bool),
Fields(Vec<JsonPath>),
Selector(PayloadSelector),
}
impl From<bool> for WithPayloadInterface {
fn from(b: bool) -> Self {
WithPayloadInterface::Bool(b)
}
}
impl Default for WithPayloadInterface {
fn default() -> Self {
WithPayloadInterface::Bool(false)
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
#[serde(untagged, rename_all = "snake_case")]
#[serde(expecting = "Expected a boolean, or an array of strings")]
pub enum WithVector {
Bool(bool),
Selector(Vec<VectorNameBuf>),
}
impl WithVector {
pub fn is_enabled(&self) -> bool {
match self {
WithVector::Bool(b) => *b,
WithVector::Selector(_) => true,
}
}
pub fn merge(&self, other: &WithVector) -> WithVector {
match (self, other) {
(WithVector::Bool(true), _) => WithVector::Bool(true),
(_, WithVector::Bool(true)) => WithVector::Bool(true),
(WithVector::Bool(false), WithVector::Bool(false)) => WithVector::Bool(false),
(WithVector::Selector(s1), WithVector::Selector(s2)) => {
WithVector::Selector(s1.iter().chain(s2).unique().cloned().collect())
}
(WithVector::Bool(false), WithVector::Selector(s)) => WithVector::Selector(s.clone()),
(WithVector::Selector(s), WithVector::Bool(false)) => WithVector::Selector(s.clone()),
}
}
}
impl From<bool> for WithVector {
fn from(b: bool) -> Self {
WithVector::Bool(b)
}
}
impl From<VectorNameBuf> for WithVector {
fn from(name: VectorNameBuf) -> Self {
WithVector::Selector(vec![name])
}
}
impl Default for WithVector {
fn default() -> Self {
WithVector::Bool(false)
}
}
impl WithPayloadInterface {
pub fn is_required(&self) -> bool {
match self {
WithPayloadInterface::Bool(b) => *b,
WithPayloadInterface::Fields(_) | WithPayloadInterface::Selector(_) => true,
}
}
}
impl From<bool> for WithPayload {
fn from(x: bool) -> Self {
WithPayload {
enable: x,
payload_selector: None,
}
}
}
impl From<WithPayloadInterface> for WithPayload {
fn from(interface: WithPayloadInterface) -> Self {
match interface {
WithPayloadInterface::Bool(enable) => WithPayload {
enable,
payload_selector: None,
},
WithPayloadInterface::Fields(fields) => WithPayload {
enable: true,
payload_selector: Some(PayloadSelector::new_include(fields)),
},
WithPayloadInterface::Selector(selector) => WithPayload {
enable: true,
payload_selector: Some(selector),
},
}
}
}
impl From<&WithPayloadInterface> for WithPayload {
fn from(interface: &WithPayloadInterface) -> Self {
WithPayload::from(interface.clone())
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
#[serde(deny_unknown_fields, rename_all = "snake_case")]
pub struct PayloadSelectorInclude {
pub include: Vec<PayloadKeyType>,
}
impl PayloadSelectorInclude {
pub fn new(include: Vec<PayloadKeyType>) -> Self {
Self { include }
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
#[serde(deny_unknown_fields, rename_all = "snake_case")]
pub struct PayloadSelectorExclude {
pub exclude: Vec<PayloadKeyType>,
}
impl PayloadSelectorExclude {
pub fn new(exclude: Vec<PayloadKeyType>) -> Self {
Self { exclude }
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)]
#[serde(untagged, rename_all = "snake_case")]
pub enum PayloadSelector {
Include(PayloadSelectorInclude),
Exclude(PayloadSelectorExclude),
}
impl From<PayloadSelectorExclude> for WithPayloadInterface {
fn from(selector: PayloadSelectorExclude) -> Self {
WithPayloadInterface::Selector(PayloadSelector::Exclude(selector))
}
}
impl From<PayloadSelectorInclude> for WithPayloadInterface {
fn from(selector: PayloadSelectorInclude) -> Self {
WithPayloadInterface::Selector(PayloadSelector::Include(selector))
}
}
impl PayloadSelector {
pub fn new_include(vecs_payload_key_type: Vec<PayloadKeyType>) -> Self {
PayloadSelector::Include(PayloadSelectorInclude {
include: vecs_payload_key_type,
})
}
pub fn new_exclude(vecs_payload_key_type: Vec<PayloadKeyType>) -> Self {
PayloadSelector::Exclude(PayloadSelectorExclude {
exclude: vecs_payload_key_type,
})
}
pub fn process(&self, x: Payload) -> Payload {
match self {
PayloadSelector::Include(selector) => JsonPath::value_filter(&x.0, |key, _| {
selector
.include
.iter()
.any(|pattern| pattern.check_include_pattern(key))
})
.into(),
PayloadSelector::Exclude(selector) => JsonPath::value_filter(&x.0, |key, _| {
selector
.exclude
.iter()
.all(|pattern| !pattern.check_exclude_pattern(key))
})
.into(),
}
}
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Default, PartialEq, Eq)]
#[serde(deny_unknown_fields, rename_all = "snake_case")]
pub struct WithPayload {
pub enable: bool,
pub payload_selector: Option<PayloadSelector>,
}
#[derive(
Debug, Deserialize, Serialize, JsonSchema, Validate, Clone, PartialEq, Eq, Default, Hash,
)]
#[serde(rename_all = "snake_case")]
pub struct MinShould {
#[validate(nested)]
pub conditions: Vec<Condition>,
pub min_count: usize,
}
#[derive(
Debug, Deserialize, Serialize, JsonSchema, Validate, Clone, PartialEq, Eq, Default, Hash,
)]
#[serde(deny_unknown_fields, rename_all = "snake_case")]
pub struct Filter {
#[validate(nested)]
#[serde(
default,
with = "MaybeOneOrMany",
skip_serializing_if = "Option::is_none"
)]
#[schemars(with = "MaybeOneOrMany<Condition>")]
pub should: Option<Vec<Condition>>,
#[validate(nested)]
#[serde(skip_serializing_if = "Option::is_none")]
pub min_should: Option<MinShould>,
#[validate(nested)]
#[serde(
default,
with = "MaybeOneOrMany",
skip_serializing_if = "Option::is_none"
)]
#[schemars(with = "MaybeOneOrMany<Condition>")]
pub must: Option<Vec<Condition>>,
#[validate(nested)]
#[serde(
default,
with = "MaybeOneOrMany",
skip_serializing_if = "Option::is_none"
)]
#[schemars(with = "MaybeOneOrMany<Condition>")]
pub must_not: Option<Vec<Condition>>,
}
impl Filter {
pub fn new() -> Self {
Filter {
should: None,
min_should: None,
must: None,
must_not: None,
}
}
pub fn new_should(condition: Condition) -> Self {
Filter {
should: Some(vec![condition]),
min_should: None,
must: None,
must_not: None,
}
}
pub fn new_min_should(min_should: MinShould) -> Self {
Filter {
should: None,
min_should: Some(min_should),
must: None,
must_not: None,
}
}
pub fn new_must(condition: Condition) -> Self {
Filter {
should: None,
min_should: None,
must: Some(vec![condition]),
must_not: None,
}
}
pub fn new_must_not(condition: Condition) -> Self {
Filter {
should: None,
min_should: None,
must: None,
must_not: Some(vec![condition]),
}
}
pub fn with_point_ids(self, ids: impl IntoIterator<Item = PointIdType>) -> Filter {
let has_id_condition: HasIdCondition = ids.into_iter().collect();
let Filter {
should,
min_should,
must,
must_not,
} = self;
let new_must = match must {
Some(mut must) => {
must.push(Condition::HasId(has_id_condition));
Some(must)
}
None => Some(vec![Condition::HasId(has_id_condition)]),
};
Filter {
should,
min_should,
must: new_must,
must_not,
}
}
pub fn merge(&self, other: &Filter) -> Filter {
self.clone().merge_owned(other.clone())
}
pub fn merge_owned(self, other: Filter) -> Filter {
let merge_component = |this, other| -> Option<Vec<Condition>> {
match (this, other) {
(None, None) => None,
(Some(this), None) => Some(this),
(None, Some(other)) => Some(other),
(Some(mut this), Some(mut other)) => {
this.append(&mut other);
Some(this)
}
}
};
Filter {
should: merge_component(self.should, other.should),
min_should: {
match (self.min_should, other.min_should) {
(None, None) => None,
(Some(this), None) => Some(this),
(None, Some(other)) => Some(other),
(Some(mut this), Some(mut other)) => {
this.conditions.append(&mut other.conditions);
this.min_count = this.min_count.max(other.min_count);
Some(this)
}
}
},
must: merge_component(self.must, other.must),
must_not: merge_component(self.must_not, other.must_not),
}
}
pub fn merge_opts(this: Option<Self>, other: Option<Self>) -> Option<Self> {
match (this, other) {
(None, None) => None,
(Some(this), None) => Some(this),
(None, Some(other)) => Some(other),
(Some(this), Some(other)) => Some(this.merge_owned(other)),
}
}
pub fn iter_conditions(&self) -> impl Iterator<Item = &Condition> {
self.must
.iter()
.flatten()
.chain(self.must_not.iter().flatten())
.chain(self.should.iter().flatten())
.chain(self.min_should.iter().flat_map(|i| &i.conditions))
}
pub fn total_conditions_count(&self) -> usize {
fn count_all_conditions(field: Option<&Vec<Condition>>) -> usize {
field
.map(|i| i.iter().map(|j| j.sub_conditions_count()).sum::<usize>())
.unwrap_or(0)
}
count_all_conditions(self.should.as_ref())
+ count_all_conditions(self.min_should.as_ref().map(|i| &i.conditions))
+ count_all_conditions(self.must.as_ref())
+ count_all_conditions(self.must_not.as_ref())
}
pub fn max_condition_input_size(&self) -> usize {
self.iter_conditions()
.map(|i| i.size_estimation())
.max()
.unwrap_or(0)
}
}
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum SnapshotFormat {
Ancient,
Regular,
Streamable,
}
#[cfg(test)]
pub(crate) mod test_utils {
use super::{GeoLineString, GeoPoint, GeoPolygon};
pub fn build_polygon(exterior_points: Vec<(f64, f64)>) -> GeoPolygon {
let exterior_line = GeoLineString {
points: exterior_points
.into_iter()
.map(|(lon, lat)| GeoPoint::new_unchecked(lon, lat))
.collect(),
};
GeoPolygon {
exterior: exterior_line,
interiors: None,
}
}
pub fn build_polygon_with_interiors(
exterior_points: Vec<(f64, f64)>,
interiors_points: Vec<Vec<(f64, f64)>>,
) -> GeoPolygon {
let exterior_line = GeoLineString {
points: exterior_points
.into_iter()
.map(|(lon, lat)| GeoPoint::new_unchecked(lon, lat))
.collect(),
};
let interior_lines = Some(
interiors_points
.into_iter()
.map(|points| GeoLineString {
points: points
.into_iter()
.map(|(lon, lat)| GeoPoint::new_unchecked(lon, lat))
.collect(),
})
.collect(),
);
GeoPolygon {
exterior: exterior_line,
interiors: interior_lines,
}
}
}
#[cfg(test)]
mod tests {
#![expect(clippy::wildcard_enum_match_arm, reason = "test code")]
use itertools::Itertools;
use rstest::rstest;
use serde_json;
use super::test_utils::build_polygon_with_interiors;
use super::*;
#[test]
#[ignore]
fn test_rmp_vs_cbor_deserialize() {
let payload = payload_json! {"payload_key": "payload_value"};
let raw = rmp_serde::to_vec(&payload).unwrap();
let de_record: Payload = serde_cbor::from_slice(&raw).unwrap();
eprintln!("payload = {payload:#?}");
eprintln!("de_record = {de_record:#?}");
}
#[rstest]
#[case::rfc_3339("2020-03-01T00:00:00Z")]
#[case::rfc_3339_custom_tz("2020-03-01T00:00:00-09:00")]
#[case::rfc_3339_custom_tz_no_colon("2020-03-01 00:00:00-0900")]
#[case::rfc_3339_custom_tz_no_colon_and_t("2020-03-01T00:00:00-0900")]
#[case::rfc_3339_custom_tz_no_minutes("2020-03-01 00:00:00-09")]
#[case::rfc_3339_and_decimals("2020-03-01T00:00:00.123456Z")]
#[case::without_z("2020-03-01T00:00:00")]
#[case::without_z_and_decimals("2020-03-01T00:00:00.12")]
#[case::space_sep_without_z("2020-03-01 00:00:00")]
#[case::space_sep_without_z_and_decimals("2020-03-01 00:00:00.123456")]
#[case::t_sep_without_seconds("2020-03-01T00:00")]
#[case::space_sep_without_seconds("2020-03-01 00:00")]
#[case::date_only("2020-03-01")]
fn test_datetime_deserialization(#[case] datetime: &str) {
let datetime = DateTimePayloadType::from_str(datetime).unwrap();
let serialized = serde_json::to_string(&datetime).unwrap();
let deserialized: DateTimePayloadType = serde_json::from_str(&serialized).unwrap();
assert_eq!(datetime, deserialized);
}
#[test]
fn test_datetime_deserialization_equivalency() {
let datetime_str = "2020-03-01T01:02:03.123456Z";
let datetime_str_no_z = "2020-03-01T01:02:03.123456";
let datetime = DateTimePayloadType::from_str(datetime_str).unwrap();
let datetime_no_z = DateTimePayloadType::from_str(datetime_str_no_z).unwrap();
assert_eq!(datetime.timestamp(), datetime_no_z.timestamp());
}
#[test]
fn test_invalid_datetime_range_returns_clear_rfc3339_error() {
let json = r#"{
"key": "created_at",
"range": {
"gte": "2014-01-01T00:00:00BAD"
}
}"#;
let err = serde_json::from_str::<Condition>(json)
.unwrap_err()
.to_string();
assert!(err.contains("RFC3339"), "err was: {err}");
assert!(err.contains("2014-01-01T00:00:00BAD"), "err was: {err}");
assert!(err.contains("Example"), "err was: {err}");
}
#[test]
fn test_datetime_payload_type_binary_roundtrip() {
let original = DateTimePayloadType::from_str("2024-06-15T12:30:45Z").unwrap();
let binary = rmp_serde::to_vec(&original).expect("serialize");
let restored: DateTimePayloadType = rmp_serde::from_slice(&binary).expect("deserialize");
assert_eq!(original, restored);
}
#[test]
fn test_range_interface_datetime_binary_roundtrip() {
let dt_gte = DateTimePayloadType::from_str("2024-01-01T00:00:00Z").unwrap();
let dt_lte = DateTimePayloadType::from_str("2024-12-31T23:59:59Z").unwrap();
let range = RangeInterface::DateTime(Range {
lt: None,
gt: None,
gte: Some(dt_gte),
lte: Some(dt_lte),
});
let binary = rmp_serde::to_vec(&range).expect("serialize");
let restored: RangeInterface = rmp_serde::from_slice(&binary).expect("deserialize");
assert_eq!(range, restored);
}
#[test]
fn test_condition_json_fallback_to_untagged() {
let is_empty_json = r#"{"is_empty": {"key": "optional_field"}}"#;
let condition: Condition = serde_json::from_str(is_empty_json).unwrap();
assert!(matches!(condition, Condition::IsEmpty(_)));
let has_id_json = r#"{"has_id": [1, 2, 3]}"#;
let condition: Condition = serde_json::from_str(has_id_json).unwrap();
assert!(matches!(condition, Condition::HasId(_)));
let nested_json = r#"{"nested": {"key": "items", "filter": {"must": []}}}"#;
let condition: Condition = serde_json::from_str(nested_json).unwrap();
assert!(matches!(condition, Condition::Nested(_)));
}
#[test]
fn test_datetime_wrapper_transcoding() {
let expected = DateTimeWrapper(chrono::Utc::now());
let transcoded = DateTimeWrapper::from_str(&expected.to_string()).unwrap();
assert_eq!(expected, transcoded);
}
#[test]
fn test_timezone_ordering() {
let datetimes = [
"2000-06-08 00:18:53+0900",
"2000-06-07 07:25:34-1100",
"2000-07-10T00:18:53+0100",
"2000-07-11 00:25:34-01:00",
"2000-07-11 00:25:35-01",
];
let sorted_datetimes: Vec<_> = datetimes
.iter()
.enumerate()
.map(|(i, s)| (i, DateTimePayloadType::from_str(s).unwrap()))
.sorted_by_key(|(_, dt)| dt.timestamp())
.collect();
sorted_datetimes
.array_windows()
.for_each(|[(i1, dt1), (i2, dt2)]| {
assert!(
i1 < i2,
"i1: {}, dt1: {}, ts1: {}\ni2: {}, dt2: {}, ts2: {}",
i1,
dt1.0,
dt1.timestamp(),
i2,
dt2.0,
dt2.timestamp()
);
});
}
#[test]
fn test_geo_radius_check_point() {
let radius = GeoRadius {
center: GeoPoint::new_unchecked(0.0, 0.0),
radius: OrderedFloat(80000.0),
};
let inside_result = radius.check_point(&GeoPoint::new_unchecked(0.5, 0.5));
assert!(inside_result);
let outside_result = radius.check_point(&GeoPoint::new_unchecked(1.5, 1.5));
assert!(!outside_result);
}
#[test]
fn test_geo_boundingbox_check_point() {
let bounding_box = GeoBoundingBox {
top_left: GeoPoint::new_unchecked(-1.0, 1.0),
bottom_right: GeoPoint::new_unchecked(1.0, -1.0),
};
let inside_result = bounding_box.check_point(&GeoPoint::new_unchecked(-0.5, 0.5));
assert!(inside_result);
let outside_result = bounding_box.check_point(&GeoPoint::new_unchecked(1.5, 1.5));
assert!(!outside_result);
}
#[test]
fn test_geo_boundingbox_antimeridian_check_point() {
let bounding_box = GeoBoundingBox {
top_left: GeoPoint::new_unchecked(167.0, 74.071028),
bottom_right: GeoPoint::new_unchecked(-66.885417, 18.7763),
};
let inside_result =
bounding_box.check_point(&GeoPoint::new_unchecked(-73.991516, 40.75798));
assert!(inside_result);
let outside_result = bounding_box.check_point(&GeoPoint::new_unchecked(13.41053, 52.52437));
assert!(!outside_result);
}
#[test]
fn test_geo_polygon_check_point() {
let test_cases = [
(
vec![
(-1.0, -1.0),
(1.0, -1.0),
(1.0, 1.0),
(-1.0, 1.0),
(-1.0, -1.0),
],
vec![vec![]],
vec![((0.5, 0.5), true), ((1.5, 1.5), false), ((1.0, 0.0), false)],
),
(
vec![
(-1.0, -1.0),
(1.0, 1.0),
(1.0, -1.0),
(-1.0, 1.0),
(-1.0, -1.0),
],
vec![vec![]],
vec![((0.5, 0.0), true), ((0.0, 0.5), false), ((0.0, 0.0), false)],
),
(
vec![
(-1.0, -1.0),
(1.5, -1.0),
(1.5, 1.5),
(-1.0, 1.5),
(-1.0, -1.0),
],
vec![vec![
(-0.5, -0.5),
(-0.5, 0.5),
(0.5, 0.5),
(0.5, -0.5),
(-0.5, -0.5),
]],
vec![((0.6, 0.6), true), ((0.0, 0.0), false), ((0.5, 0.5), false)],
),
];
for (exterior, interiors, points) in test_cases {
let polygon = build_polygon_with_interiors(exterior, interiors);
for ((lon, lat), expected_result) in points {
let inside_result = polygon
.convert()
.check_point(&GeoPoint::new_unchecked(lon, lat));
assert_eq!(inside_result, expected_result);
}
}
}
#[test]
fn test_serialize_query() {
let filter = Filter {
must: Some(vec![Condition::Field(FieldCondition::new_match(
JsonPath::new("hello"),
"world".to_owned().into(),
))]),
must_not: None,
should: None,
min_should: None,
};
let json = serde_json::to_string_pretty(&filter).unwrap();
eprintln!("{json}")
}
#[test]
fn test_deny_unknown_fields() {
let query1 = r#"
{
"wrong": "query"
}
"#;
let filter: Result<Filter, _> = serde_json::from_str(query1);
assert!(filter.is_err())
}
#[test]
fn test_parse_match_query() {
let query = r#"
{
"key": "hello",
"match": { "value": 42 }
}
"#;
let condition: FieldCondition = serde_json::from_str(query).unwrap();
assert_eq!(
condition.r#match.unwrap(),
Match::Value(MatchValue {
value: ValueVariants::Integer(42)
})
);
let query = r#"
{
"key": "hello",
"match": { "value": true }
}
"#;
let condition: FieldCondition = serde_json::from_str(query).unwrap();
assert_eq!(
condition.r#match.unwrap(),
Match::Value(MatchValue {
value: ValueVariants::Bool(true)
})
);
let query = r#"
{
"key": "hello",
"match": { "value": "world" }
}
"#;
let condition: FieldCondition = serde_json::from_str(query).unwrap();
assert_eq!(
condition.r#match.unwrap(),
Match::Value(MatchValue {
value: ValueVariants::String("world".to_owned())
})
);
}
#[test]
fn test_parse_match_any() {
let query = r#"
{
"should": [
{
"key": "Jason",
"match": {
"any": [
"Bourne",
"Momoa",
"Statham"
]
}
}
]
}
"#;
let filter: Filter = serde_json::from_str(query).unwrap();
let should = filter.should.unwrap();
assert_eq!(should.len(), 1);
let Some(Condition::Field(c)) = should.first() else {
panic!("Condition::Field expected")
};
assert_eq!(c.key.to_string(), "Jason");
let Match::Any(m) = c.r#match.as_ref().unwrap() else {
panic!("Match::Any expected")
};
if let AnyVariants::Strings(kws) = &m.any {
assert_eq!(kws.len(), 3);
let expect: IndexSet<_, FnvBuildHasher> = ["Bourne", "Momoa", "Statham"]
.into_iter()
.map(|i| i.to_string())
.collect();
assert_eq!(kws, &expect);
} else {
panic!("AnyVariants::Keywords expected");
}
}
#[test]
fn test_parse_match_any_mixed_types() {
let query = r#"
{
"should": [
{
"key": "Jason",
"match": {
"any": [
"Bourne",
42
]
}
}
]
}
"#;
let result: Result<Filter, _> = serde_json::from_str(query);
assert!(result.is_err());
}
#[test]
fn test_parse_nested_match_query() {
let query = r#"
{
"key": "hello.nested",
"match": { "value": 42 }
}
"#;
let condition: FieldCondition = serde_json::from_str(query).unwrap();
assert_eq!(
condition.r#match.unwrap(),
Match::Value(MatchValue {
value: ValueVariants::Integer(42)
})
);
let query = r#"
{
"key": "hello.nested",
"match": { "value": true }
}
"#;
let condition: FieldCondition = serde_json::from_str(query).unwrap();
assert_eq!(
condition.r#match.unwrap(),
Match::Value(MatchValue {
value: ValueVariants::Bool(true)
})
);
let query = r#"
{
"key": "hello.nested",
"match": { "value": "world" }
}
"#;
let condition: FieldCondition = serde_json::from_str(query).unwrap();
assert_eq!(
condition.r#match.unwrap(),
Match::Value(MatchValue {
value: ValueVariants::String("world".to_owned())
})
);
}
#[test]
fn test_parse_empty_query() {
let query = r#"
{
"should": [
{
"is_empty" : {
"key" : "Jason"
}
}
]
}
"#;
let filter: Filter = serde_json::from_str(query).unwrap();
let should = filter.should.unwrap();
assert_eq!(should.len(), 1);
let Some(Condition::IsEmpty(c)) = should.first() else {
panic!("Condition::IsEmpty expected")
};
assert_eq!(c.is_empty.key.to_string(), "Jason");
}
#[test]
fn test_parse_null_query() {
let query = r#"
{
"should": [
{
"is_null" : {
"key" : "Jason"
}
}
]
}
"#;
let filter: Filter = serde_json::from_str(query).unwrap();
let should = filter.should.unwrap();
assert_eq!(should.len(), 1);
let Some(Condition::IsNull(c)) = should.first() else {
panic!("Condition::IsNull expected")
};
assert_eq!(c.is_null.key.to_string(), "Jason");
}
#[test]
fn test_parse_nested_filter_query() {
let query = r#"
{
"must": [
{
"nested": {
"key": "country.cities",
"filter": {
"must": [
{
"key": "population",
"range": {
"gte": 8
}
},
{
"key": "sightseeing",
"values_count": {
"lt": 3
}
}
]
}
}
}
]
}
"#;
let filter: Filter = serde_json::from_str(query).unwrap();
let musts = filter.must.unwrap();
assert_eq!(musts.len(), 1);
match musts.first() {
Some(Condition::Nested(nested_condition)) => {
assert_eq!(nested_condition.raw_key().to_string(), "country.cities");
assert_eq!(nested_condition.array_key().to_string(), "country.cities[]");
let nested_musts = nested_condition.filter().must.as_ref().unwrap();
assert_eq!(nested_musts.len(), 2);
let first_must = nested_musts.first().unwrap();
match first_must {
Condition::Field(c) => {
assert_eq!(c.key.to_string(), "population");
assert!(c.range.is_some());
}
_ => panic!("Condition::Field expected"),
}
let second_must = nested_musts.get(1).unwrap();
match second_must {
Condition::Field(c) => {
assert_eq!(c.key.to_string(), "sightseeing");
assert!(c.values_count.is_some());
}
_ => panic!("Condition::Field expected"),
}
}
o => panic!("Condition::Nested expected but got {o:?}"),
};
}
#[test]
fn test_parse_single_nested_filter_query() {
let query = r#"
{
"must": {
"nested": {
"key": "country.cities",
"filter": {
"must": {
"key": "population",
"range": {
"gte": 8
}
}
}
}
}
}
"#;
let filter: Filter = serde_json::from_str(query).unwrap();
let musts = filter.must.unwrap();
assert_eq!(musts.len(), 1);
let first_must = musts.first().unwrap();
let Condition::Nested(nested_condition) = first_must else {
panic!("Condition::Nested expected but got {first_must:?}")
};
assert_eq!(nested_condition.raw_key().to_string(), "country.cities");
assert_eq!(nested_condition.array_key().to_string(), "country.cities[]");
let nested_must = nested_condition.filter().must.as_ref().unwrap();
assert_eq!(nested_must.len(), 1);
let must = nested_must.first().unwrap();
let Condition::Field(c) = must else {
panic!("Condition::Field expected, got {must:?}")
};
assert_eq!(c.key.to_string(), "population");
assert!(c.range.is_some());
}
#[test]
fn test_payload_query_parse() {
let query1 = r#"
{
"must": [
{
"key": "hello",
"match": {
"value": 42
}
},
{
"must_not": [
{
"has_id": [1, 2, 3, 4]
},
{
"key": "geo_field",
"geo_bounding_box": {
"top_left": {
"lon": 13.410146,
"lat": 52.519289
},
"bottom_right": {
"lon": 13.432683,
"lat": 52.505582
}
}
}
]
}
]
}
"#;
let filter: Filter = serde_json::from_str(query1).unwrap();
eprintln!("{filter:?}");
let must = filter.must.unwrap();
let _must_not = filter.must_not;
assert_eq!(must.len(), 2);
match must.get(1) {
Some(Condition::Filter(f)) => {
let must_not = &f.must_not;
match must_not {
Some(v) => assert_eq!(v.len(), 2),
None => panic!("Filter expected"),
}
}
_ => panic!("Condition expected"),
}
}
#[test]
fn test_nested_payload_query_parse() {
let query1 = r#"
{
"must": [
{
"key": "hello.nested.world",
"match": {
"value": 42
}
},
{
"key": "foo.nested.bar",
"match": {
"value": 1
}
}
]
}
"#;
let filter: Filter = serde_json::from_str(query1).unwrap();
let must = filter.must.unwrap();
assert_eq!(must.len(), 2);
}
#[test]
fn test_min_should_query_parse() {
let query1 = r#"
{
"min_should": {
"conditions": [
{
"key": "hello.nested.world",
"match": {
"value": 42
}
},
{
"key": "foo.nested.bar",
"match": {
"value": 1
}
}
],
"min_count": 2
}
}
"#;
let filter: Filter = serde_json::from_str(query1).unwrap();
let min_should = filter.min_should.unwrap();
assert_eq!(min_should.conditions.len(), 2);
}
#[test]
fn test_min_should_nested_parse() {
let query1 = r#"
{
"must": [
{
"min_should": {
"conditions": [
{
"key": "hello.nested.world",
"match": {
"value": 42
}
},
{
"key": "foo.nested.bar",
"match": {
"value": 1
}
}
],
"min_count": 2
}
}
]
}
"#;
let filter: Filter = serde_json::from_str(query1).unwrap();
let must = filter.must.unwrap();
assert_eq!(must.len(), 1);
match must.first() {
Some(Condition::Filter(f)) => {
let min_should = &f.min_should;
match min_should {
Some(v) => assert_eq!(v.conditions.len(), 2),
None => panic!("Filter expected"),
}
}
_ => panic!("Condition expected"),
}
}
#[test]
fn test_geo_validation() {
let query1 = r#"
{
"must": [
{
"key": "geo_field",
"geo_bounding_box": {
"top_left": {
"lon": 1113.410146,
"lat": 52.519289
},
"bottom_right": {
"lon": 13.432683,
"lat": 52.505582
}
}
}
]
}
"#;
let filter: Result<Filter, _> = serde_json::from_str(query1);
assert!(filter.is_err());
let query2 = r#"
{
"must": [
{
"key": "geo_field",
"geo_polygon": {
"exterior": {},
"interiors": []
}
}
]
}
"#;
let filter: Result<Filter, _> = serde_json::from_str(query2);
assert!(filter.is_err());
let query3 = r#"
{
"must": [
{
"key": "geo_field",
"geo_polygon": {
"exterior":{
"points": [
{"lon": -12.0, "lat": -34.0},
{"lon": 11.0, "lat": -22.0},
{"lon": -32.0, "lat": -14.0}
]
},
"interiors": []
}
}
]
}
"#;
let filter: Result<Filter, _> = serde_json::from_str(query3);
assert!(filter.is_err());
let query4 = r#"
{
"must": [
{
"key": "geo_field",
"geo_polygon": {
"exterior": {
"points": [
{"lon": -12.0, "lat": -34.0},
{"lon": 11.0, "lat": -22.0},
{"lon": -32.0, "lat": -14.0},
{"lon": -12.0, "lat": -34.0}
]
},
"interiors": []
}
}
]
}
"#;
let filter: Result<Filter, _> = serde_json::from_str(query4);
assert!(filter.is_ok());
let query5 = r#"
{
"must": [
{
"key": "geo_field",
"geo_polygon": {
"exterior": {
"points": [
{"lon": -12.0, "lat": -34.0},
{"lon": 11.0, "lat": -22.0},
{"lon": -32.0, "lat": -14.0},
{"lon": -12.0, "lat": -34.0}
]
},
"interiors": [
{
"points": [
{"lon": -12.0, "lat": -34.0},
{"lon": 11.0, "lat": -22.0},
{"lon": -32.0, "lat": -14.0}
]
}
]
}
}
]
}
"#;
let filter: Result<Filter, _> = serde_json::from_str(query5);
assert!(filter.is_err());
let query6 = r#"
{
"must": [
{
"key": "geo_field",
"geo_polygon": {
"exterior": {
"points": [
{"lon": -12.0, "lat": -34.0},
{"lon": 11.0, "lat": -22.0},
{"lon": -32.0, "lat": -14.0},
{"lon": -12.0, "lat": -34.0}
]
},
"interiors": [
{
"points": [
{"lon": -12.0, "lat": -34.0},
{"lon": 11.0, "lat": -22.0},
{"lon": -32.0, "lat": -14.0},
{"lon": -12.0, "lat": -34.0}
]
}
]
}
}
]
}
"#;
let filter: Result<Filter, _> = serde_json::from_str(query6);
assert!(filter.is_ok());
}
#[test]
fn test_payload_parsing() {
let ft = PayloadFieldSchema::FieldType(PayloadSchemaType::Keyword);
let ft_json = serde_json::to_string(&ft).unwrap();
eprintln!("ft_json = {ft_json:?}");
let ft = PayloadFieldSchema::FieldParams(PayloadSchemaParams::Text(Default::default()));
let ft_json = serde_json::to_string(&ft).unwrap();
eprintln!("ft_json = {ft_json:?}");
let query = r#""keyword""#;
let field_type: PayloadSchemaType = serde_json::from_str(query).unwrap();
eprintln!("field_type = {field_type:?}");
}
#[test]
fn merge_filters() {
let condition1 = Condition::Field(FieldCondition::new_match(
JsonPath::new("summary"),
Match::new_text("Berlin"),
));
let mut this = Filter::new_must(condition1.clone());
this.should = Some(vec![condition1.clone()]);
let condition2 = Condition::Field(FieldCondition::new_match(
JsonPath::new("city"),
Match::new_value(ValueVariants::String("Osaka".into())),
));
let other = Filter::new_must(condition2.clone());
let merged = this.merge(&other);
assert!(merged.must.is_some());
assert_eq!(merged.must.as_ref().unwrap().len(), 2);
assert!(merged.must_not.is_none());
assert!(merged.should.is_some());
assert_eq!(merged.should.as_ref().unwrap().len(), 1);
assert!(merged.must.as_ref().unwrap().contains(&condition1));
assert!(merged.must.as_ref().unwrap().contains(&condition2));
assert!(merged.should.as_ref().unwrap().contains(&condition1));
}
#[test]
fn test_payload_selector_include() {
let payload = payload_json! {
"a": 1,
"b": {
"c": 123,
"e": {
"f": [1,2,3],
"g": 7,
"h": "text",
"i": [
{
"j": 1,
"k": 2
},
{
"j": 3,
"k": 4
}
]
}
}
};
let selector =
PayloadSelector::new_include(vec![JsonPath::new("a"), JsonPath::new("b.e.f")]);
let payload = selector.process(payload);
let expected = payload_json! {
"a": 1,
"b": {
"e": {
"f": [1,2,3],
}
}
};
assert_eq!(payload, expected);
}
#[test]
fn test_payload_selector_array_include() {
let payload = payload_json! {
"a": 1,
"b": {
"c": 123,
"f": [1,2,3,4,5],
}
};
let selector = PayloadSelector::new_include(vec![JsonPath::new("a"), JsonPath::new("a")]);
let payload = selector.process(payload);
let expected = payload_json! {
"a": 1
};
assert_eq!(payload, expected);
let selector = PayloadSelector::new_include(vec![JsonPath::new("b.f[0]")]);
let payload = selector.process(payload);
let expected = payload_json! {};
assert_eq!(payload, expected);
}
#[test]
fn test_payload_selector_no_implicit_array_include() {
let payload = payload_json! {
"a": 1,
"b": {
"c": [
{
"d": 1,
"e": 2
},
{
"d": 3,
"e": 4
}
],
}
};
let selector = PayloadSelector::new_include(vec![JsonPath::new("b.c")]);
let selected_payload = selector.process(payload.clone());
let expected = payload_json! {
"b": {
"c": [
{
"d": 1,
"e": 2
},
{
"d": 3,
"e": 4
}
]
}
};
assert_eq!(selected_payload, expected);
let selector = PayloadSelector::new_include(vec![JsonPath::new("b.c[].d")]);
let selected_payload = selector.process(payload.clone());
let expected = payload_json! {
"b": {
"c": [
{"d": 1},
{"d": 3}
]
}
};
assert_eq!(selected_payload, expected);
let selector = PayloadSelector::new_include(vec![JsonPath::new("b.c.d")]);
let selected_payload = selector.process(payload);
let expected = payload_json! {
"b": {
"c": []
}
};
assert_eq!(selected_payload, expected);
}
#[test]
fn test_payload_selector_exclude() {
let payload = payload_json! {
"a": 1,
"b": {
"c": 123,
"e": {
"f": [1,2,3],
"g": 7,
"h": "text",
"i": [
{
"j": 1,
"k": 2
},
{
"j": 3,
"k": 4
}
]
}
}
};
let selector =
PayloadSelector::new_exclude(vec![JsonPath::new("a"), JsonPath::new("b.e.f")]);
let payload = selector.process(payload);
let expected = payload_json! {
"b": {
"c": 123,
"e": {
"g": 7,
"h": "text",
"i": [
{
"j": 1,
"k": 2
},
{
"j": 3,
"k": 4
}
]
}
}
};
assert_eq!(payload, expected);
}
#[test]
fn test_payload_selector_array_exclude() {
let payload = payload_json! {
"a": 1,
"b": {
"c": 123,
"f": [1,2,3,4,5],
}
};
let selector = PayloadSelector::new_exclude(vec![JsonPath::new("a"), JsonPath::new("a")]);
let payload = selector.process(payload);
let expected = payload_json! {
"b": {
"c": 123,
"f": [1,2,3,4,5],
}
};
assert_eq!(payload, expected);
let selector = PayloadSelector::new_exclude(vec![JsonPath::new("b.f[0]")]);
let payload = selector.process(payload);
let expected = payload_json! {
"b": {
"c": 123,
"f": [1,2,3,4,5],
}
};
assert_eq!(payload, expected);
}
#[test]
fn test_extended_point_id_cbor_roundtrip() {
let uuid = Uuid::parse_str("550e8400-e29b-41d4-a716-446655440000").unwrap();
for point_id in [ExtendedPointId::Uuid(uuid), ExtendedPointId::NumId(42)] {
let cbor_bytes = serde_cbor::to_vec(&point_id).unwrap();
let deserialized: ExtendedPointId = serde_cbor::from_slice(&cbor_bytes).unwrap();
assert_eq!(point_id, deserialized);
}
}
#[test]
fn test_filter_with_match_and_has_id_uuid_cbor_roundtrip() {
let uuid = Uuid::parse_str("550e8400-e29b-41d4-a716-446655440000").unwrap();
let filter = Filter {
should: None,
min_should: None,
must: Some(vec![Condition::Field(FieldCondition::new_match(
crate::segment::json_path::JsonPath::new("org_id"),
Match::new_value(ValueVariants::String("test_org".to_string())),
))]),
must_not: Some(vec![Condition::HasId(HasIdCondition {
has_id: [ExtendedPointId::Uuid(uuid)].into_iter().collect(),
})]),
};
let cbor_bytes = serde_cbor::to_vec(&filter).unwrap();
let deserialized: Filter = serde_cbor::from_slice(&cbor_bytes).unwrap();
assert_eq!(filter, deserialized);
}
}
fn shard_key_string_example() -> String {
"region_1".to_string()
}
fn shard_key_number_example() -> u64 {
12
}
#[derive(Deserialize, Serialize, JsonSchema, Debug, Clone, PartialEq, Eq, Hash)]
#[serde(untagged)]
pub enum ShardKey {
#[schemars(
schema_with = "String::json_schema",
example = "shard_key_string_example"
)]
Keyword(EcoString),
#[schemars(example = "shard_key_number_example")]
Number(u64),
}
impl From<String> for ShardKey {
fn from(s: String) -> Self {
ShardKey::Keyword(EcoString::from(s))
}
}
impl From<EcoString> for ShardKey {
fn from(s: EcoString) -> Self {
ShardKey::Keyword(s)
}
}
impl From<&str> for ShardKey {
fn from(s: &str) -> Self {
ShardKey::Keyword(EcoString::from(s))
}
}
impl From<u64> for ShardKey {
fn from(n: u64) -> Self {
ShardKey::Number(n)
}
}
impl Display for ShardKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ShardKey::Keyword(keyword) => write!(f, "\"{keyword}\""),
ShardKey::Number(number) => write!(f, "{number}"),
}
}
}