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

common_query/
stream.rs
