qdrant/lib/edge/python/src/lib.rs

pub mod bm25;
pub mod config;
pub mod count;
pub mod facet;
pub mod info;
pub mod query;
pub mod repr;
pub mod scroll;
pub mod search;
pub mod snapshots;
pub mod types;
pub mod update;
pub mod utils;

use std::path::PathBuf;

use bytemuck::TransparentWrapperAlloc as _;
use edge::EdgeConfig;
use pyo3::exceptions::PyException;
use pyo3::prelude::*;
use segment::common::operation_error::OperationError;
use segment::types::*;

use self::config::*;
use self::count::*;
use self::facet::*;
use self::info::*;
use self::query::*;
use self::scroll::*;
use self::search::*;
use self::types::*;
use self::update::*;

#[pymodule]
mod qdrant_edge {
    #[pymodule_export]
    use super::PyEdgeShard;
    #[pymodule_export]
    use super::bm25::{PyBm25, PyBm25Config};
    #[pymodule_export]
    use super::config::quantization::{
        PyBinaryQuantizationConfig, PyBinaryQuantizationEncoding,
        PyBinaryQuantizationQueryEncoding, PyCompressionRatio, PyProductQuantizationConfig,
        PyScalarQuantizationConfig, PyScalarType, PyTurboQuantBitSize,
        PyTurboQuantQuantizationConfig,
    };
    #[pymodule_export]
    use super::config::sparse_vector_data::{PyEdgeSparseVectorParams, PyModifier};
    #[pymodule_export]
    use super::config::vector_data::{
        PyDistance, PyEdgeVectorParams, PyHnswIndexConfig, PyMultiVectorComparator,
        PyMultiVectorConfig, PyPlainIndexConfig, PyVectorStorageDatatype,
    };
    #[pymodule_export]
    use super::config::{PyEdgeConfig, PyEdgeOptimizersConfig};
    #[pymodule_export]
    use super::count::PyCountRequest;
    #[pymodule_export]
    use super::facet::{PyFacetHit, PyFacetRequest, PyFacetResponse};
    #[pymodule_export]
    use super::query::{
        PyDirection, PyFusion, PyMmr, PyOrderBy, PyPrefetch, PyQueryRequest, PySample,
    };
    #[pymodule_export]
    use super::scroll::PyScrollRequest;
    #[pymodule_export]
    use super::search::{
        PyAcornSearchParams, PyQuantizationSearchParams, PySearchParams, PySearchRequest,
    };
    #[pymodule_export]
    use super::types::filter::{
        PyFieldCondition, PyFilter, PyGeoBoundingBox, PyGeoPoint, PyGeoPolygon, PyGeoRadius,
        PyHasIdCondition, PyHasVectorCondition, PyIsEmptyCondition, PyIsNullCondition, PyMatchAny,
        PyMatchExcept, PyMatchPhrase, PyMatchText, PyMatchTextAny, PyMatchValue, PyMinShould,
        PyNestedCondition, PyRangeDateTime, PyRangeFloat, PyValuesCount,
    };
    #[pymodule_export]
    use super::types::formula::{PyDecayKind, PyExpressionInterface, PyFormula};
    #[pymodule_export]
    use super::types::payload_schema::{
        PyBoolIndexParams, PyDatetimeIndexParams, PyFloatIndexParams, PyGeoIndexParams,
        PyIntegerIndexParams, PyKeywordIndexParams, PyLanguage, PyPayloadSchemaType,
        PySnowballLanguage, PySnowballParams, PyStopwordsSet, PyTextIndexParams, PyTokenizerType,
        PyUuidIndexParams,
    };
    #[pymodule_export]
    use super::types::query::{
        PyContextPair, PyContextQuery, PyDiscoverQuery, PyFeedbackItem, PyFeedbackNaiveQuery,
        PyNaiveFeedbackCoefficients, PyPayloadSelectorInterface, PyQueryInterface,
        PyRecommendQuery,
    };
    #[pymodule_export]
    use super::types::{PyPoint, PyPointVectors, PyRecord, PyScoredPoint, PySparseVector};
    #[pymodule_export]
    use super::update::{PyUpdateMode, PyUpdateOperation};
}

#[pyclass(name = "EdgeShard")]
#[derive(Debug)]
pub struct PyEdgeShard(Option<edge::EdgeShard>);

#[pymethods]
impl PyEdgeShard {
    /// Load an edge shard from existing files at `path`.
    /// Optional `config`: if provided, compatibility is checked and config is overwritten on disk.
    #[staticmethod]
    #[pyo3(signature = (path, config = None))]
    pub fn load(path: PathBuf, config: Option<PyEdgeConfig>) -> Result<Self> {
        let shard = edge::EdgeShard::load(&path, config.map(EdgeConfig::from))?;
        Ok(Self(Some(shard)))
    }

    /// Create a new edge shard at `path` with the given configuration.
    /// Fails if the path already contains segment data.
    #[staticmethod]
    pub fn create(path: PathBuf, config: PyEdgeConfig) -> Result<Self> {
        let shard = edge::EdgeShard::new(&path, config.0)?;
        Ok(Self(Some(shard)))
    }

    pub fn flush(&self) -> Result<()> {
        self.get_shard()?.flush();
        Ok(())
    }

    pub fn optimize(&self) -> Result<bool> {
        let optimized = self.get_shard()?.optimize()?;
        Ok(optimized)
    }

    pub fn close(&mut self) {
        self.0.take(); // `edge::Shard` is automatically flushed on drop
    }

    pub fn update(&self, operation: PyUpdateOperation) -> Result<()> {
        self.get_shard()?.update(operation.into())?;
        Ok(())
    }

    pub fn query(&self, query: PyQueryRequest) -> Result<Vec<PyScoredPoint>> {
        let points = self.get_shard()?.query(query.into())?;
        let points = PyScoredPoint::wrap_vec(points);
        Ok(points)
    }

    pub fn search(&self, search: PySearchRequest) -> Result<Vec<PyScoredPoint>> {
        let points = self.get_shard()?.search(search.into())?;
        let points = PyScoredPoint::wrap_vec(points);
        Ok(points)
    }

    pub fn scroll(&self, scroll: PyScrollRequest) -> Result<(Vec<PyRecord>, Option<PyPointId>)> {
        let (points, next_offset) = self.get_shard()?.scroll(scroll.into())?;
        let points = PyRecord::wrap_vec(points);
        Ok((points, next_offset.map(PyPointId)))
    }

    pub fn count(&self, count: PyCountRequest) -> Result<usize> {
        let points_count = self.get_shard()?.count(count.into())?;
        Ok(points_count)
    }

    pub fn facet(&self, facet: PyFacetRequest) -> Result<PyFacetResponse> {
        let response = self.get_shard()?.facet(facet.into())?;
        Ok(PyFacetResponse::new(response))
    }

    pub fn retrieve(
        &self,
        point_ids: Vec<PyPointId>,
        with_payload: Option<PyWithPayload>,
        with_vector: Option<PyWithVector>,
    ) -> Result<Vec<PyRecord>> {
        let point_ids = PyPointId::peel_vec(point_ids);
        let points = self.get_shard()?.retrieve(
            &point_ids,
            with_payload.map(WithPayloadInterface::from),
            with_vector.map(WithVector::from),
        )?;
        let points = PyRecord::wrap_vec(points);
        Ok(points)
    }

    pub fn info(&self) -> Result<PyShardInfo> {
        let info = self.get_shard()?.info();
        let info = PyShardInfo(info);
        Ok(info)
    }

    // ------- Snapshot related methods -------

    #[staticmethod]
    pub fn unpack_snapshot(snapshot_path: PathBuf, target_path: PathBuf) -> Result<()> {
        edge::EdgeShard::unpack_snapshot(&snapshot_path, &target_path)?;
        Ok(())
    }

    pub fn snapshot_manifest(&self) -> Result<PyValue> {
        let manifest = self.get_shard()?.snapshot_manifest()?;
        Ok(PyValue::new(serde_json::to_value(&manifest).unwrap()))
    }

    #[pyo3(signature = (snapshot_path, tmp_dir=None))]
    pub fn update_from_snapshot(
        &mut self,
        snapshot_path: PathBuf,
        tmp_dir: Option<PathBuf>,
    ) -> Result<()> {
        self._update_from_snapshot(snapshot_path, tmp_dir)?;
        Ok(())
    }
}

pub type Result<T, E = PyError> = std::result::Result<T, E>;

#[derive(Debug)]
pub struct PyError(OperationError);

impl From<OperationError> for PyError {
    fn from(err: OperationError) -> Self {
        Self(err)
    }
}

impl From<PyError> for PyErr {
    fn from(err: PyError) -> Self {
        PyException::new_err(err.0.to_string())
    }
}