qdrant_datafusion/
stream.rs

1//! Generate a stream of `Qdrant` points as arrow `RecordBatch`.
2
3use std::pin::Pin;
4use std::sync::Arc;
5use std::task::{Context, Poll};
6
7use datafusion::arrow::array::RecordBatch;
8use datafusion::arrow::datatypes::*;
9use datafusion::error::Result as DataFusionResult;
10use datafusion::execution::RecordBatchStream;
11use futures_util::Stream;
12
13/// Stream that yields `RecordBatch`es from Qdrant query results.
14///
15/// This stream implementation provides the bridge between `Qdrant`'s async query results
16/// and `DataFusion`'s streaming execution model. It typically yields a single batch
17/// containing all points from a `Qdrant` query, though it's designed to be extensible
18/// for future pagination support.
19///
20/// # Implementation Notes
21/// Currently optimized for `Qdrant`'s typical usage patterns where queries return
22/// relatively small result sets in a single response. Future versions may add
23/// support for streaming large result sets with pagination.
24#[pin_project::pin_project]
25pub struct QdrantQueryStream {
26    schema: SchemaRef,
27    #[pin]
28    stream: Pin<Box<dyn Stream<Item = DataFusionResult<RecordBatch>> + Send>>,
29}
30
31impl QdrantQueryStream {
32    /// Create a new stream that yields record batches with the specified schema.
33    ///
34    /// # Arguments
35    /// * `schema` - The Arrow schema that defines the structure of record batches
36    /// * `stream` - The underlying stream of record batch results
37    ///
38    /// # Returns
39    /// A new `QdrantQueryStream` ready for `DataFusion` execution.
40    pub fn new(
41        schema: SchemaRef,
42        stream: Pin<Box<dyn Stream<Item = DataFusionResult<RecordBatch>> + Send>>,
43    ) -> Self {
44        Self { schema, stream }
45    }
46}
47
48impl Stream for QdrantQueryStream {
49    type Item = datafusion::error::Result<RecordBatch>;
50
51    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
52        self.as_mut().project().stream.poll_next(cx)
53    }
54}
55
56impl RecordBatchStream for QdrantQueryStream {
57    fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) }
58}
qdrant_datafusion/stream.rs

qdrant_datafusion/
stream.rs
