qdrant-datafusion 0.1.1

Qdrant integration for Apache DataFusion
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

//! Various utility functions for working with schema and data.

use std::collections::HashSet;

use datafusion::arrow::datatypes::Schema;

/// Specification for selecting which vectors to retrieve from `Qdrant`.
///
/// This enum determines what vector data will be requested from `Qdrant` during query execution.
/// The selection is optimized based on the `DataFusion` schema projection.
#[derive(Debug, Clone)]
pub enum VectorSelectorSpec {
    /// No vectors needed - only metadata fields (id, payload) requested
    None,
    /// All vectors needed - either unnamed collection or all named vectors requested
    All,
    /// Specific named vectors needed - only fetch these vector fields
    Named(Vec<String>),
}

/// Build an optimal vector selector based on the projected schema.
///
/// This function analyzes the `DataFusion` schema (after projection) to determine
/// which vector fields are actually needed, enabling efficient queries that only
/// fetch required data from `Qdrant`.
///
/// # Arguments
/// * `schema` - The Arrow schema (potentially projected) defining which fields are needed
///
/// # Returns
/// A `VectorSelectorSpec` that tells `Qdrant` exactly which vectors to include in the response.
///
/// # Examples
/// ```rust,ignore
/// use datafusion::arrow::datatypes::{Schema, Field, DataType};
/// use qdrant_datafusion::utils::{build_vector_selector, VectorSelectorSpec};
/// use std::sync::Arc;
///
/// // Schema with only metadata - no vectors needed
/// let metadata_schema = Schema::new(vec![
///     Field::new("id", DataType::Utf8, false),
///     Field::new("payload", DataType::Utf8, true),
/// ]);
/// assert!(matches!(build_vector_selector(&metadata_schema), VectorSelectorSpec::None));
///
/// // Schema with unnamed vector - fetch all
/// let unnamed_schema = Schema::new(vec![
///     Field::new("id", DataType::Utf8, false),
///     Field::new("vector", DataType::List(Arc::new(Field::new("item", DataType::Float32, true))), true),
/// ]);
/// assert!(matches!(build_vector_selector(&unnamed_schema), VectorSelectorSpec::All));
/// ```
pub fn build_vector_selector(schema: &Schema) -> VectorSelectorSpec {
    let mut vector_names: HashSet<_> = schema
        .fields()
        .iter()
        .filter(|f| !["id", "payload"].contains(&f.name().as_str()))
        .map(|f| f.name())
        .map(|name| {
            if name.ends_with("_indices") || name.ends_with("_values") {
                // Extract base name for sparse vectors
                name.trim_end_matches("_indices").trim_end_matches("_values")
            } else {
                name
            }
        })
        .collect();

    if vector_names.is_empty() {
        return VectorSelectorSpec::None;
    }

    let has_unnamed_vector = vector_names.remove("vector");
    if has_unnamed_vector || vector_names.is_empty() {
        // Unnamed vector collection - use simple boolean selector
        VectorSelectorSpec::All
    } else {
        // Specific named vectors
        VectorSelectorSpec::Named(
            vector_names.into_iter().map(ToString::to_string).collect::<Vec<_>>(),
        )
    }
}

/// Determine if payload data should be included in `Qdrant` queries.
///
/// This function checks if the schema includes a "payload" field, which indicates
/// that payload data is needed and should be fetched from `Qdrant`.
///
/// # Arguments
/// * `schema` - The Arrow schema to analyze
///
/// # Returns
/// `true` if payload field is present in schema, `false` otherwise
///
/// # Examples
/// ```rust,ignore
/// use datafusion::arrow::datatypes::{Schema, Field, DataType};
/// use qdrant_datafusion::utils::build_payload_selector;
///
/// // Schema with payload field
/// let with_payload = Schema::new(vec![
///     Field::new("id", DataType::Utf8, false),
///     Field::new("payload", DataType::Utf8, true),
/// ]);
/// assert_eq!(build_payload_selector(&with_payload), true);
///
/// // Schema without payload
/// let no_payload = Schema::new(vec![
///     Field::new("id", DataType::Utf8, false),
/// ]);
/// assert_eq!(build_payload_selector(&no_payload), false);
/// ```
pub fn build_payload_selector(schema: &Schema) -> bool {
    schema.fields().iter().any(|f| f.name() == "payload")
}