qdrant_datafusion/
table.rs

1//! `DataFusion` `TableProvider` implementation for `Qdrant` vector database collections.
2use std::any::Any;
3use std::sync::Arc;
4
5use datafusion::arrow::array::*;
6use datafusion::arrow::datatypes::*;
7use datafusion::catalog::{Session, TableProvider};
8use datafusion::datasource::TableType;
9use datafusion::error::{DataFusionError, Result as DataFusionResult};
10use datafusion::execution::{SendableRecordBatchStream, TaskContext};
11use datafusion::logical_expr::dml::InsertOp;
12use datafusion::physical_plan::execution_plan::Boundedness;
13use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
14use datafusion::prelude::Expr;
15use datafusion::sql::TableReference;
16use qdrant_client::Qdrant;
17use qdrant_client::qdrant::{QueryPointsBuilder, VectorsSelector};
18
19use crate::arrow::deserialize::QdrantRecordBatchBuilder;
20use crate::arrow::schema::collection_to_arrow_schema;
21use crate::error::{Error, Result};
22use crate::stream::QdrantQueryStream;
23use crate::utils;
24
25/// `DataFusion` `TableProvider` implementation for `Qdrant` vector database collections.
26///
27/// This is the main interface for integrating `Qdrant` collections with `DataFusion` SQL queries.
28/// It provides a complete SQL interface over vector data with support for all `Qdrant` vector
29/// types, schema projection optimization, and heterogeneous collection handling.
30///
31/// # Features
32/// - **Complete Vector Support**: Dense, multi-dense, and sparse vectors
33/// - **Schema Projection**: Only fetches vector fields that are actually requested
34/// - **Heterogeneous Collections**: Handles points with different vector field subsets
35/// - **High Performance**: Single-pass processing with minimal allocations
36///
37/// # Examples
38///
39/// ## Basic Usage
40/// ```rust,ignore
41/// use qdrant_datafusion::prelude::*;
42/// use qdrant_client::Qdrant;
43/// use datafusion::prelude::*;
44/// use std::sync::Arc;
45///
46/// # async fn example() -> Result<()> {
47/// // Connect to Qdrant
48/// let client = Qdrant::from_url("http://localhost:6334").build()?;
49///
50/// // Create table provider for a collection
51/// let table_provider = QdrantTableProvider::try_new(client, "my_vectors").await?;
52///
53/// // Register with DataFusion
54/// let ctx = SessionContext::new();
55/// ctx.register_table("vectors", Arc::new(table_provider))?;
56///
57/// // Query with SQL
58/// let df = ctx.sql("SELECT id, embedding FROM vectors LIMIT 10").await?;
59/// let results = df.collect().await?;
60/// # Ok(())
61/// # }
62/// ```
63///
64/// ## Advanced Projections
65/// ```rust,no_run
66/// # use qdrant_datafusion::prelude::*;
67/// # use datafusion::prelude::*;
68/// # async fn example(ctx: SessionContext) -> Result<()> {
69/// // Only fetch specific vector fields (optimized query to Qdrant)
70/// let df = ctx.sql("
71///     SELECT
72///         text_embedding,
73///         keywords_indices,
74///         keywords_values
75///     FROM mixed_vectors
76///     WHERE id = 'doc123'
77/// ").await?;
78/// # Ok(())
79/// # }
80/// ```
81#[derive(Clone)]
82pub struct QdrantTableProvider {
83    table:  TableReference,
84    client: Arc<Qdrant>,
85    schema: Arc<Schema>,
86}
87
88impl std::fmt::Debug for QdrantTableProvider {
89    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
90        f.debug_struct("QdrantTableProvider")
91            .field("table", &self.table)
92            .field("client", &"Qdrant")
93            .field("schema", &self.schema)
94            .finish()
95    }
96}
97
98impl QdrantTableProvider {
99    /// Create a new `QdrantTableProvider` for the specified `Qdrant` collection.
100    ///
101    /// This constructor connects to the `Qdrant` collection, analyzes its schema, and creates
102    /// a DataFusion-compatible table provider. The schema is built by examining the collection's
103    /// vector configuration and creating appropriate Arrow fields for all vector types.
104    ///
105    /// # Arguments
106    /// * `client` - Connected `Qdrant` client instance
107    /// * `collection` - Name of the `Qdrant` collection to provide access to
108    ///
109    /// # Returns
110    /// A configured `QdrantTableProvider` ready for SQL queries.
111    ///
112    /// # Errors
113    /// Returns an error if:
114    /// - The collection does not exist or is inaccessible
115    /// - The collection configuration cannot be retrieved
116    /// - The collection has an unsupported schema configuration
117    ///
118    /// # Examples
119    /// ```rust,ignore
120    /// use qdrant_datafusion::prelude::*;
121    /// use qdrant_client::Qdrant;
122    ///
123    /// # async fn example() -> Result<()> {
124    /// let client = Qdrant::from_url("http://localhost:6334")
125    ///     .api_key("optional-api-key")
126    ///     .build()?;
127    ///
128    /// let table_provider = QdrantTableProvider::try_new(client, "embeddings").await?;
129    /// # Ok(())
130    /// # }
131    /// ```
132    pub async fn try_new(client: Qdrant, collection: &str) -> Result<Self> {
133        let info = client.collection_info(collection).await?;
134        // Get the config
135        let config = info
136            .result
137            .ok_or(Error::MissingCollectionInfo(collection.into()))?
138            .config
139            .ok_or(Error::MissingCollectionInfo(collection.into()))?;
140        let schema = collection_to_arrow_schema(collection, &config)?;
141        Ok(Self {
142            table:  TableReference::bare(collection),
143            client: Arc::new(client),
144            schema: Arc::new(schema),
145        })
146    }
147}
148
149#[async_trait::async_trait]
150impl TableProvider for QdrantTableProvider {
151    fn as_any(&self) -> &dyn Any { self }
152
153    fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) }
154
155    fn table_type(&self) -> TableType { TableType::Base }
156
157    async fn scan(
158        &self,
159        _state: &dyn Session,
160        projection: Option<&Vec<usize>>,
161        filters: &[Expr],
162        limit: Option<usize>,
163    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
164        // Apply projection to schema ONCE, here
165        let projected_schema = match projection {
166            Some(indices) if !indices.is_empty() => Arc::new(self.schema.project(indices)?),
167            _ => Arc::clone(&self.schema),
168        };
169
170        // Build selectors based on what fields are in the projected schema
171        let vector_selector = utils::build_vector_selector(&projected_schema);
172        let payload_selector = utils::build_payload_selector(&projected_schema);
173
174        // For now, ignore filters - we'll handle them with UDFs later
175        Ok(Arc::new(QdrantScanExec::new(
176            Arc::clone(&self.client),
177            self.table.table().to_string(),
178            projected_schema,
179            vector_selector,
180            payload_selector,
181            filters,
182            limit,
183        )))
184    }
185
186    async fn insert_into(
187        &self,
188        _state: &dyn Session,
189        _input: Arc<dyn ExecutionPlan>,
190        _insert_op: InsertOp,
191    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
192        todo!()
193    }
194}
195
196/// `DataFusion` `ExecutionPlan` implementation for scanning Qdrant collections.
197///
198/// This is the physical execution plan node that actually performs queries against `Qdrant`.
199/// It's created by the `QdrantTableProvider` during query planning and handles the execution
200/// of `Qdrant` queries with optimizations like vector field selection and payload filtering.
201///
202/// # Features
203/// - **Optimized Vector Selection**: Only fetches vector fields that are needed
204/// - **Schema Projection**: Respects `DataFusion` column pruning
205/// - **Async Streaming**: Non-blocking execution with proper backpressure
206/// - **Limit Pushdown**: Limit constraints are pushed to Qdrant for efficiency
207///
208/// This struct is typically not used directly - it's created automatically by the
209/// `QdrantTableProvider` during SQL query execution.
210#[derive(Clone)]
211pub struct QdrantScanExec {
212    client:           Arc<Qdrant>,
213    collection:       String,
214    schema:           SchemaRef, // Already projected
215    vector_selector:  utils::VectorSelectorSpec,
216    payload_selector: bool,
217    filter:           Arc<[Expr]>,
218    limit:            Option<usize>,
219    properties:       PlanProperties,
220}
221
222impl std::fmt::Debug for QdrantScanExec {
223    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
224        f.debug_struct("QdrantScanExec")
225            .field("client", &"Qdrant")
226            .field("collection", &self.collection)
227            .field("schema", &self.schema)
228            .field("vector_selector", &self.vector_selector)
229            .field("payload_selector", &self.payload_selector)
230            .field("limit", &self.limit)
231            .finish_non_exhaustive()
232    }
233}
234
235impl QdrantScanExec {
236    pub fn new(
237        client: Arc<Qdrant>,
238        collection: String,
239        schema: SchemaRef,
240        vector_selector: utils::VectorSelectorSpec,
241        payload_selector: bool,
242        filter: &[Expr],
243        limit: Option<usize>,
244    ) -> Self {
245        let properties = PlanProperties::new(
246            datafusion::physical_expr::EquivalenceProperties::new(Arc::clone(&schema)),
247            datafusion::physical_plan::Partitioning::UnknownPartitioning(1),
248            datafusion::physical_plan::execution_plan::EmissionType::Final,
249            Boundedness::Bounded,
250        );
251
252        Self {
253            client,
254            collection,
255            schema,
256            vector_selector,
257            payload_selector,
258            filter: Arc::from(filter),
259            limit,
260            properties,
261        }
262    }
263}
264
265/// Execute a `Qdrant` query and return a `RecordBatch`.
266///
267/// # Errors
268/// - Returns an error if the query fails.
269pub(crate) async fn execute_qdrant_query(
270    client: Arc<Qdrant>,
271    collection: String,
272    schema: SchemaRef,
273    vector_selector: utils::VectorSelectorSpec,
274    payload_selector: bool,
275    _filters: &[Expr],
276    limit: Option<usize>,
277) -> DataFusionResult<RecordBatch> {
278    // Build query using QueryPointsBuilder
279    let mut query_builder = QueryPointsBuilder::new(&collection);
280
281    // Use the builder's API which accepts Into<SelectorOptions>
282    match vector_selector {
283        utils::VectorSelectorSpec::None => {
284            query_builder = query_builder.with_vectors(false);
285        }
286        utils::VectorSelectorSpec::All => {
287            query_builder = query_builder.with_vectors(true);
288        }
289        utils::VectorSelectorSpec::Named(names) => {
290            query_builder = query_builder.with_vectors(VectorsSelector { names });
291        }
292    }
293
294    query_builder = query_builder.with_payload(payload_selector);
295
296    if let Some(limit_val) = limit {
297        query_builder = query_builder.limit(limit_val as u64);
298    }
299
300    // Execute query
301    let response =
302        client.query(query_builder).await.map_err(|e| DataFusionError::External(Box::new(e)))?;
303
304    // Convert points to RecordBatch using incremental builder
305    let points = response.result;
306
307    if points.is_empty() {
308        return Ok(RecordBatch::new_empty(schema));
309    }
310
311    // Create incremental builder with pre-allocated capacity
312    let mut builder = QdrantRecordBatchBuilder::new(schema, points.len());
313
314    // Single pass through points with true owned iteration
315    for point in points {
316        builder.append_point(point); // Pass owned point, not borrowed
317    }
318
319    builder.finish()
320}
321
322impl ExecutionPlan for QdrantScanExec {
323    fn name(&self) -> &'static str { "QdrantScanExec" }
324
325    fn as_any(&self) -> &dyn Any { self }
326
327    fn properties(&self) -> &PlanProperties { &self.properties }
328
329    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { vec![] }
330
331    fn with_new_children(
332        self: Arc<Self>,
333        _children: Vec<Arc<dyn ExecutionPlan>>,
334    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
335        Ok(self)
336    }
337
338    fn execute(
339        &self,
340        _partition: usize,
341        _context: Arc<TaskContext>,
342    ) -> DataFusionResult<SendableRecordBatchStream> {
343        let client = Arc::clone(&self.client);
344        let collection = self.collection.clone();
345        let schema = Arc::clone(&self.schema);
346        let vector_selector = self.vector_selector.clone();
347        let payload_selector = self.payload_selector;
348        let filter = Arc::clone(&self.filter);
349        let limit = self.limit;
350        let inner = Box::pin(futures_util::stream::once(async move {
351            execute_qdrant_query(
352                client,
353                collection,
354                schema,
355                vector_selector,
356                payload_selector,
357                &filter,
358                limit,
359            )
360            .await
361        }));
362        let stream = QdrantQueryStream::new(Arc::clone(&self.schema), inner);
363        Ok(Box::pin(stream))
364    }
365}
366
367impl DisplayAs for QdrantScanExec {
368    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
369        match t {
370            DisplayFormatType::Default | DisplayFormatType::Verbose => {
371                write!(f, "QdrantScanExec: collection={}", self.collection)?;
372                if let Some(limit) = self.limit {
373                    write!(f, ", limit={limit}")?;
374                }
375                Ok(())
376            }
377            DisplayFormatType::TreeRender => {
378                write!(f, "QdrantScanExec: collection={}", self.collection)
379            }
380        }
381    }
382}
qdrant_datafusion/table.rs

qdrant_datafusion/
table.rs
