common_query/
stream.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::any::Any;
16use std::fmt::{Debug, Formatter};
17use std::sync::{Arc, Mutex};
18
19use common_recordbatch::SendableRecordBatchStream;
20use common_recordbatch::adapter::DfRecordBatchStreamAdapter;
21use datafusion::execution::SendableRecordBatchStream as DfSendableRecordBatchStream;
22use datafusion::execution::context::TaskContext;
23use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalSortExpr};
24use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
25use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
26use datafusion_common::DataFusionError;
27use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef;
28use datatypes::schema::SchemaRef;
29
30/// Opens a new scan stream equivalent to the one this plan was built from.
31pub type StreamFactoryRef =
32    Arc<dyn Fn() -> datafusion_common::Result<SendableRecordBatchStream> + Send + Sync>;
33
34/// Adapts greptime's [SendableRecordBatchStream] to DataFusion's [ExecutionPlan].
35pub struct StreamScanAdapter {
36    stream: Mutex<Option<SendableRecordBatchStream>>,
37    /// Opens a replacement stream once `stream` has been handed out. DataFusion
38    /// executes a plan more than once in a recursive CTE, where every iteration
39    /// re-executes the recursive term.
40    stream_factory: Option<StreamFactoryRef>,
41    schema: SchemaRef,
42    arrow_schema: ArrowSchemaRef,
43    properties: Arc<PlanProperties>,
44    output_ordering: Option<Vec<PhysicalSortExpr>>,
45}
46
47impl Debug for StreamScanAdapter {
48    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
49        f.debug_struct("StreamScanAdapter")
50            .field("stream", &"<SendableRecordBatchStream>")
51            .field("schema", &self.schema)
52            .finish()
53    }
54}
55
56impl StreamScanAdapter {
57    pub fn new(stream: SendableRecordBatchStream) -> Self {
58        let schema = stream.schema();
59        let arrow_schema = schema.arrow_schema().clone();
60        let properties = Arc::new(PlanProperties::new(
61            EquivalenceProperties::new(arrow_schema.clone()),
62            Partitioning::UnknownPartitioning(1),
63            EmissionType::Incremental,
64            Boundedness::Bounded,
65        ));
66
67        Self {
68            stream: Mutex::new(Some(stream)),
69            stream_factory: None,
70            schema,
71            arrow_schema,
72            properties,
73            output_ordering: None,
74        }
75    }
76
77    pub fn with_output_ordering(mut self, output_ordering: Option<Vec<PhysicalSortExpr>>) -> Self {
78        self.output_ordering = output_ordering;
79        self
80    }
81
82    /// Makes this plan re-executable. The factory must open a stream over the
83    /// same scan request.
84    pub fn with_stream_factory(mut self, stream_factory: StreamFactoryRef) -> Self {
85        self.stream_factory = Some(stream_factory);
86        self
87    }
88}
89
90impl DisplayAs for StreamScanAdapter {
91    fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
92        write!(
93            f,
94            "StreamScanAdapter: [<SendableRecordBatchStream>], schema: ["
95        )?;
96        write!(f, "{:?}", &self.arrow_schema)?;
97        write!(f, "]")
98    }
99}
100
101impl ExecutionPlan for StreamScanAdapter {
102    fn as_any(&self) -> &dyn Any {
103        self
104    }
105
106    fn schema(&self) -> ArrowSchemaRef {
107        self.arrow_schema.clone()
108    }
109
110    fn properties(&self) -> &Arc<PlanProperties> {
111        &self.properties
112    }
113
114    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
115        vec![]
116    }
117
118    // DataFusion will swap children unconditionally.
119    // But since this node is leaf node, it's safe to just return self.
120    fn with_new_children(
121        self: Arc<Self>,
122        _children: Vec<Arc<dyn ExecutionPlan>>,
123    ) -> datafusion_common::Result<Arc<dyn ExecutionPlan>> {
124        Ok(self.clone())
125    }
126
127    fn execute(
128        &self,
129        _partition: usize,
130        _context: Arc<TaskContext>,
131    ) -> datafusion_common::Result<DfSendableRecordBatchStream> {
132        let stream = self.stream.lock().unwrap().take();
133        let stream = match stream {
134            Some(stream) => stream,
135            None => {
136                let factory = self.stream_factory.as_ref().ok_or_else(|| {
137                    DataFusionError::Execution("Stream already exhausted".to_string())
138                })?;
139                factory()?
140            }
141        };
142        Ok(Box::pin(DfRecordBatchStreamAdapter::new(stream)))
143    }
144
145    fn name(&self) -> &str {
146        "StreamScanAdapter"
147    }
148}
149
150#[cfg(test)]
151mod test {
152    use common_recordbatch::{RecordBatch, RecordBatches};
153    use datafusion::prelude::SessionContext;
154    use datatypes::data_type::ConcreteDataType;
155    use datatypes::schema::{ColumnSchema, Schema};
156    use datatypes::vectors::Int32Vector;
157    use futures_util::TryStreamExt;
158
159    use super::*;
160
161    #[tokio::test]
162    async fn test_simple_table_scan() {
163        let ctx = SessionContext::new();
164        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
165            "a",
166            ConcreteDataType::int32_datatype(),
167            false,
168        )]));
169
170        let batch1 = RecordBatch::new(
171            schema.clone(),
172            vec![Arc::new(Int32Vector::from_slice([1, 2])) as _],
173        )
174        .unwrap();
175        let batch2 = RecordBatch::new(
176            schema.clone(),
177            vec![Arc::new(Int32Vector::from_slice([3, 4, 5])) as _],
178        )
179        .unwrap();
180
181        let recordbatches =
182            RecordBatches::try_new(schema.clone(), vec![batch1.clone(), batch2.clone()]).unwrap();
183        let stream = recordbatches.as_stream();
184
185        let scan = StreamScanAdapter::new(stream);
186
187        assert_eq!(scan.schema(), schema.arrow_schema().clone());
188
189        let stream = scan.execute(0, ctx.task_ctx()).unwrap();
190        let recordbatches = stream.try_collect::<Vec<_>>().await.unwrap();
191        assert_eq!(recordbatches[0], batch1.into_df_record_batch());
192        assert_eq!(recordbatches[1], batch2.into_df_record_batch());
193
194        let result = scan.execute(0, ctx.task_ctx());
195        assert!(result.is_err());
196        match result {
197            Err(e) => assert!(e.to_string().contains("Stream already exhausted")),
198            _ => unreachable!(),
199        }
200    }
201
202    #[tokio::test]
203    async fn test_re_execute_with_stream_factory() {
204        let ctx = SessionContext::new();
205        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
206            "a",
207            ConcreteDataType::int32_datatype(),
208            false,
209        )]));
210
211        let batch = RecordBatch::new(
212            schema.clone(),
213            vec![Arc::new(Int32Vector::from_slice([1, 2])) as _],
214        )
215        .unwrap();
216
217        let factory_schema = schema.clone();
218        let factory_batch = batch.clone();
219        let scan = StreamScanAdapter::new(
220            RecordBatches::try_new(schema.clone(), vec![batch.clone()])
221                .unwrap()
222                .as_stream(),
223        )
224        .with_stream_factory(Arc::new(move || {
225            Ok(
226                RecordBatches::try_new(factory_schema.clone(), vec![factory_batch.clone()])
227                    .unwrap()
228                    .as_stream(),
229            )
230        }));
231
232        for _ in 0..3 {
233            let stream = scan.execute(0, ctx.task_ctx()).unwrap();
234            let recordbatches = stream.try_collect::<Vec<_>>().await.unwrap();
235            assert_eq!(recordbatches, vec![batch.clone().into_df_record_batch()]);
236        }
237    }
238}
common_query/stream.rs

common_query/
stream.rs
