qdrant_datafusion/arrow/
schema.rs

1//! Schema utilities for `Qdrant` `DataFusion` integration.
2use std::sync::Arc;
3
4use datafusion::arrow::datatypes::*;
5use qdrant_client::qdrant::{CollectionConfig, Datatype, VectorParams, vectors_config};
6
7use crate::error::{Error, Result};
8
9/// Simple helper function to determine if a field is a multi-vector
10pub fn is_multi_vector_field(field: &Field) -> bool {
11    matches!(
12        field.data_type(),
13        DataType::List(inner) if matches!(inner.data_type(), DataType::List(_))
14    )
15}
16
17/// Simple helper function to convert a Qdrant datatype to an Arrow datatype.
18pub fn datatype_to_arrow(_dt: Datatype) -> DataType {
19    // TODO: Decide whether to support other vector data types, since Qdrant currently only ever
20    // sends f32
21    DataType::Float32
22    // match dt {
23    //     Datatype::Default | Datatype::Float32 => DataType::Float32,
24    //     Datatype::Float16 => DataType::Float16,
25    //     Datatype::Uint8 => DataType::UInt8,
26    // }
27}
28
29/// Simple helper function to create a vector field
30pub fn create_vector_field(name: &str, dt: Datatype, nullable: bool) -> FieldRef {
31    Field::new(name, datatype_to_arrow(dt), nullable).into()
32}
33
34/// Simple helper function to create a list field for vector parameters
35pub fn create_vector_param_field(name: &str, vector_params: &VectorParams) -> Field {
36    if vector_params.multivector_config.is_some() {
37        Field::new(
38            name,
39            DataType::List(Arc::new(Field::new(
40                "item",
41                DataType::List(create_vector_field("item", vector_params.datatype(), true)),
42                true,
43            ))),
44            true, // Allow nulls - points may not have all vectors
45        )
46    } else {
47        Field::new(
48            name,
49            DataType::List(create_vector_field("item", vector_params.datatype(), true)),
50            true, // Allow nulls - points may not have all vectors
51        )
52    }
53}
54
55/// Convert a collection's configuration info into an Arrow schema.
56///
57/// # Errors
58/// - Returns an error if the collection info or the vector params is missing.
59pub fn collection_to_arrow_schema(collection: &str, config: &CollectionConfig) -> Result<Schema> {
60    let mut fields = vec![
61        // The point ID (can be numeric or UUID string)
62        Field::new("id", DataType::Utf8, false),
63        // Payload as JSON string
64        Field::new("payload", DataType::Utf8, true),
65    ];
66
67    // Get the params from config
68    let params =
69        config.params.as_ref().ok_or(Error::MissingCollectionInfoParams(collection.into()))?;
70
71    // Parse vectors_config if present
72    if let Some(config) = params.vectors_config.as_ref().and_then(|c| c.config.as_ref()) {
73        match config {
74            vectors_config::Config::Params(vector_params) => {
75                // Single unnamed vector
76                fields.push(create_vector_param_field("vector", vector_params));
77            }
78            vectors_config::Config::ParamsMap(params_map) => {
79                // Multiple named vectors
80                fields.extend(
81                    params_map
82                        .map
83                        .iter()
84                        .map(|(name, params)| create_vector_param_field(name, params)),
85                );
86            }
87        }
88    }
89
90    // Parse sparse_vectors_config if present
91    if let Some(sparse_config) = &params.sparse_vectors_config {
92        // SparseVectorConfig has a map field
93        for name in sparse_config.map.keys() {
94            // Sparse indices are always u32, regardless of index config datatype
95            fields.push(Field::new(
96                format!("{name}_indices"),
97                DataType::List(Field::new("item", DataType::UInt32, true).into()),
98                true, // Allow nulls - points may not have all sparse vectors
99            ));
100            // Sparse values are always f32
101            fields.push(Field::new(
102                format!("{name}_values"),
103                DataType::List(create_vector_field("item", Datatype::Float32, true)),
104                true, // Allow nulls - points may not have all sparse vectors
105            ));
106        }
107    }
108
109    Ok(Schema::new(fields))
110}
qdrant_datafusion/arrow/schema.rs

qdrant_datafusion/arrow/
schema.rs
